Update test_einsum.py

lettucecfd · Sep 11, 2024 · ac3d10c · ac3d10c
1 parent d00c4c0
commit ac3d10c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/flow/test_einsum.py b/tests/flow/test_einsum.py
@@ -170,7 +170,7 @@ def source_term(self, u):
 
 @pytest.mark.parametrize("fix_dim", [1, 2, 3])
 def test_einsum(fix_dim):
-    context = Context()
+    context = Context(dtype=torch.float64)
 
     flow = EinsumFlow(context, [16] * fix_dim)
     flow.f = torch.rand_like(flow.f)
@@ -200,7 +200,7 @@ def test_einsum(fix_dim):
     f_log = -torch.log(torch.einsum("q...,q...->q...", [flow.f,
                                                         1 / flow.torch_stencil.w]))
     f2 = torch.einsum("q...,q...->...", [flow.f, f_log])
-    assert torch.allclose(f0, f1)
+    assert torch.allclose(f0, f1, rtol=)
     assert torch.allclose(f1, f2)
 
     f0 = flow.pseudo_entropy_global()