Use the build in function

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