[BugFix] Better return_log_prob=True for tensordict outputs

ghstack-source-id: 977af3880f39cb341c1c715f1b8c9d59b7c580a0 Pull Request resolved: #1155
pytorch · Dec 20, 2024 · b229c59 · b229c59
1 parent 863ba10
commit b229c59
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
@@ -510,9 +510,12 @@ def forward(
                         kwargs = {"aggregate_probabilities": False}
                     log_prob = dist.log_prob(out_tensors, **kwargs)
                     if log_prob is not out_tensors:
-                        # Composite dists return the tensordict_out directly when aggrgate_prob is False
-                        out_tensors.set(self.log_prob_key, log_prob)
-                    else:
+                        if is_tensor_collection(log_prob):
+                            out_tensors.update(log_prob)
+                        else:
+                            # Composite dists return the tensordict_out directly when aggrgate_prob is False
+                            out_tensors.set(self.log_prob_key, log_prob)
+                    elif dist.log_prob_key in out_tensors:
                         out_tensors.rename_key_(dist.log_prob_key, self.log_prob_key)
                 tensordict_out.update(out_tensors)
             else: