[BugFix] Robust torch.cuda.is_current_stream_capturing calls

ghstack-source-id: 8335587
Pull-Request-resolved: #2950

Author: Vincent Moens

torchrl/__init__.py

@@ -50,10 +50,13 @@
     auto_unwrap_transformed_env,
     compile_with_warmup,
     implement_for,
+    logger,
     set_auto_unwrap_transformed_env,
     timeit,
 )
 
+torchrl_logger = logger
+
 # Filter warnings in subprocesses: True by default given the multiple optional
 # deps of the library. This can be turned on via `torchrl.filter_warnings_subprocess = False`.
 filter_warnings_subprocess = True
@@ -108,4 +111,6 @@ def _inv(self):
     "implement_for",
     "set_auto_unwrap_transformed_env",
     "timeit",
+    "logger",
+    "torchrl_logger",
 ]

torchrl/_utils.py

@@ -1122,3 +1122,17 @@ def auto_unwrap_transformed_env(allow_none=False):
     elif _AUTO_UNWRAP is None:
         return _DEFAULT_AUTO_UNWRAP
     return strtobool(_AUTO_UNWRAP) if isinstance(_AUTO_UNWRAP, str) else _AUTO_UNWRAP
+
+
+def safe_is_current_stream_capturing():
+    """A safe proxy to torch.cuda.is_current_stream_capturing."""
+    if not torch.cuda.is_available():
+        return False
+    try:
+        return torch.cuda.is_current_stream_capturing()
+    except Exception as error:
+        warnings.warn(
+            f"torch.cuda.is_current_stream_capturing() exited unexpectedly with the error message {error=}. "
+            f"Returning False by default."
+        )
+        return False

torchrl/modules/distributions/continuous.py

@@ -15,6 +15,7 @@
 from torch.distributions import constraints
 from torch.distributions.transforms import _InverseTransform
 
+from torchrl._utils import safe_is_current_stream_capturing
 from torchrl.modules.distributions.truncated_normal import (
     TruncatedNormal as _TruncatedNormal,
 )
@@ -358,7 +359,7 @@ def __init__(
             event_dims = min(1, loc.ndim)
 
         err_msg = "TanhNormal high values must be strictly greater than low values"
-        if not is_compiling() and not torch.cuda.is_current_stream_capturing():
+        if not is_compiling() and not safe_is_current_stream_capturing():
             if isinstance(high, torch.Tensor) or isinstance(low, torch.Tensor):
                 if not (high > low).all():
                     raise RuntimeError(err_msg)
@@ -377,7 +378,7 @@ def __init__(
             low = torch.as_tensor(low, device=loc.device)
         elif low.device != loc.device:
             low = low.to(loc.device)
-        if not is_compiling() and not torch.cuda.is_current_stream_capturing():
+        if not is_compiling() and not safe_is_current_stream_capturing():
             self.non_trivial_max = (high != 1.0).any()
             self.non_trivial_min = (low != -1.0).any()
         else: