[BugFix] replace assert_allclose with assert_close (#644)

pytorch · Nov 4, 2022 · ada0fcd · ada0fcd
1 parent b4b27fe
commit ada0fcd
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/test/test_rb.py b/test/test_rb.py
@@ -247,16 +247,12 @@ def test_prototype_prb(priority_key, contiguous, device):
     rb.update_tensordict_priority(s)
     s, _ = rb.sample(5)
     assert (val == s.get("a")).sum() >= 1
-    torch.testing.assert_allclose(
-        td2[idx0].get("a").view(1), s.get("a").unique().view(1)
-    )
+    torch.testing.assert_close(td2[idx0].get("a").view(1), s.get("a").unique().view(1))
 
     # test updating values of original td
     td2.set_("a", torch.ones_like(td2.get("a")))
     s, _ = rb.sample(5)
-    torch.testing.assert_allclose(
-        td2[idx0].get("a").view(1), s.get("a").unique().view(1)
-    )
+    torch.testing.assert_close(td2[idx0].get("a").view(1), s.get("a").unique().view(1))
 
 
 @pytest.mark.parametrize("stack", [False, True])