fix loss test case for batch size variation (Lightning-Universe#402)

chris-clem · Dec 9, 2020 · a4d7e01 · a4d7e01
1 parent 1bd9ae7
commit a4d7e01
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/models/rl/unit/test_reinforce.py b/tests/models/rl/unit/test_reinforce.py
@@ -37,9 +37,9 @@ def setUp(self) -> None:
     def test_loss(self):
         """Test the reinforce loss function"""
 
-        batch_states = torch.rand(32, 4)
-        batch_actions = torch.rand(32).long()
-        batch_qvals = torch.rand(32)
+        batch_states = torch.rand(16, 4)
+        batch_actions = torch.rand(16).long()
+        batch_qvals = torch.rand(16)
 
         loss = self.model.loss(batch_states, batch_actions, batch_qvals)