Fix: make torch.equal torch.all(torch.isclose(...

alexhernandezgarcia · Dec 26, 2023 · bb09c7d · bb09c7d
1 parent f3180e6
commit bb09c7d
Showing 1 changed file with 10 additions and 6 deletions.
diff --git a/tests/gflownet/utils/test_batch.py b/tests/gflownet/utils/test_batch.py
@@ -1354,12 +1354,16 @@ def test__get_rewards_parents_multiple_env_returns_expected_non_terminating(
     rewards_batch = batch.get_rewards(do_non_terminating=True)
     rewards = torch.stack(rewards)
 
-    assert torch.equal(
-        rewards_parents_batch,
-        tfloat(rewards_parents, device=batch.device, float_type=batch.float),
+    assert torch.all(
+        torch.isclose(
+            rewards_parents_batch,
+            tfloat(rewards_parents, device=batch.device, float_type=batch.float),
+        )
     ), (rewards_parents, rewards_parents_batch)
 
-    assert torch.equal(
-        rewards_batch,
-        tfloat(rewards, device=batch.device, float_type=batch.float),
+    assert torch.all(
+        torch.isclose(
+            rewards_batch,
+            tfloat(rewards, device=batch.device, float_type=batch.float),
+        )
     ), (rewards, rewards_batch)