diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py index 13d02af2b36..a563f60837d 100644 --- a/torchrl/envs/common.py +++ b/torchrl/envs/common.py @@ -3266,7 +3266,7 @@ def maybe_reset(self, tensordict: TensorDictBase) -> TensorDictBase: """ if self._simple_done: done = tensordict._get_str("done", default=None) - any_done = done.any() + any_done = done is not None and done.any() if any_done: tensordict._set_str( "_reset",