amend

pytorch · May 28, 2024 · ee2bab5 · ee2bab5
1 parent db34b71
commit ee2bab5
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/torchrl/envs/libs/gym.py b/torchrl/envs/libs/gym.py
@@ -17,7 +17,7 @@
 import torch
 from packaging import version
 
-from tensordict import TensorDictBase, TensorDict
+from tensordict import TensorDict, TensorDictBase
 from torch.utils._pytree import tree_map
 
 from torchrl._utils import implement_for
@@ -1561,9 +1561,7 @@ def replace_none(nparray):
             if is_none.any():
                 # Then it is a final observation and we delegate the registration to the appropriate reader
                 nz = (~is_none).nonzero()[0][0]
-                zero_like = tree_map(
-                    lambda x: np.zeros_like(x), nparray[nz]
-                )
+                zero_like = tree_map(lambda x: np.zeros_like(x), nparray[nz])
                 for idx in is_none.nonzero()[0]:
                     nparray[idx] = zero_like
             return tree_map(lambda *x: np.stack(x), *nparray)