No need to unsqueeze

pytorch · Nov 6, 2024 · 1f33881 · 1f33881
1 parent ef794c7
commit 1f33881
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/objectives/ppo.py b/torchrl/objectives/ppo.py
@@ -463,8 +463,8 @@ def reset(self) -> None:
     def get_entropy_bonus(self, dist: d.Distribution) -> torch.Tensor:
         try:
             entropy = dist.entropy()
-            if is_tensor_collection(entropy) and hasattr(dist, "entropy_key"):
-                entropy = entropy.get(dist.entropy_key).unsqueeze(-1)
+            if is_tensor_collection(entropy):
+                entropy = entropy.get(dist.entropy_key)
         except NotImplementedError:
             x = dist.rsample((self.samples_mc_entropy,))
             log_prob = dist.log_prob(x)