[BugFix] Fix deterministic fallback when the dist has no support (#830)

pytorch · Jun 24, 2024 · 8522bb6 · 8522bb6
1 parent aa2d0bc
commit 8522bb6
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
@@ -472,9 +472,15 @@ def _dist_sample(
             try:
                 return dist.deterministic_sample
             except AttributeError:
-                fallback = (
-                    "mean" if isinstance(dist.support, D.constraints._Real) else "mode"
-                )
+                try:
+                    support = dist.support
+                    fallback = (
+                        "mean" if isinstance(support, D.constraints._Real) else "mode"
+                    )
+                except NotImplementedError:
+                    # Some custom dists don't have a support
+                    # We arbitrarily fall onto 'mean' in these cases
+                    fallback = "mean"
                 try:
                     if fallback == "mean":
                         return dist.mean