fix tests

lucidrains · Jan 22, 2025 · 872e7ba · 872e7ba
2 parents daf6807 + f4a263f
commit 872e7ba
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/test_titans.py b/tests/test_titans.py
@@ -8,6 +8,9 @@
 def exists(v):
     return v is not None
 
+def diff(x, y):
+    return (x - y).abs().amax()
+
 @pytest.mark.parametrize('seq_len', (32, 1024, 77))
 @pytest.mark.parametrize('silu', (False, True))
 @pytest.mark.parametrize('learned_mem_model_weights', (False, True))
@@ -164,7 +167,7 @@ def test_neural_mem_inference(
 
     sequential_retrieved = torch.cat(sequential_retrieved, dim = -2)
 
-    assert torch.allclose(parallel_retrieved, sequential_retrieved, atol = 1e-6)
+    assert torch.allclose(parallel_retrieved, sequential_retrieved, atol = 1e-5)
 
 @pytest.mark.parametrize('seq_len', (1023, 17))
 @pytest.mark.parametrize('sliding', (True, False))