fixed end of batch size mismatch (#389)

Co-authored-by: unknown <[email protected]>
Lightning-Universe · Nov 22, 2020 · f0e2bee · f0e2bee
1 parent cff6b52
commit f0e2bee
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pl_bolts/models/rl/reinforce_model.py b/pl_bolts/models/rl/reinforce_model.py
@@ -215,7 +215,7 @@ def loss(self, states, actions, scaled_rewards) -> torch.Tensor:
 
         # policy loss
         log_prob = log_softmax(logits, dim=1)
-        log_prob_actions = scaled_rewards * log_prob[range(self.batch_size), actions]
+        log_prob_actions = scaled_rewards * log_prob[range(len(log_prob)), actions]
         loss = -log_prob_actions.mean()
 
         return loss