Fix test

pytorch · Jan 10, 2025 · d9a4dab · d9a4dab
1 parent f93e330
commit d9a4dab
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -12630,9 +12630,7 @@ def test_transform_env(self, num_rewards, weights):
 
         expected = sum(
             w * r
-            for w, r in zip(
-                weights, rollout.get(("next", "reward")).split(1, dim=-1), strict=True
-            )
+            for w, r in zip(weights, rollout.get(("next", "reward")).split(1, dim=-1))
         )
         torch.testing.assert_close(scalar_reward, expected)
 
@@ -12654,7 +12652,7 @@ def test_transform_model(self, num_rewards, weights):
         td = TensorDict({"reward": torch.randn(num_rewards)}, [])
         model(td)
 
-        expected = sum(w * r for w, r in zip(weights, td["reward"], strict=True))
+        expected = sum(w * r for w, r in zip(weights, td["reward"]))
         torch.testing.assert_close(td["scalar_reward"], expected)
 
     @pytest.mark.parametrize("rbclass", [ReplayBuffer, TensorDictReplayBuffer])