Fixing flaky inference test

mlfoundations · Jul 3, 2024 · 5373a56 · 5373a56
1 parent 6a5068e
commit 5373a56
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_inference_simple.py b/tests/test_inference_simple.py
@@ -48,4 +48,4 @@ def test_inference_simple(
 
         text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
 
-    assert text_probs.cpu().numpy()[0].tolist() == [1.0, 0.0, 0.0]
+    assert torch.allclose(text_probs.cpu()[0], torch.tensor([1.0, 0.0, 0.0]))