fixed end of batch size mismatch

Lightning-Universe · Nov 19, 2020 · 4e8ad0c · 4e8ad0c
1 parent fa5a944
commit 4e8ad0c
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