amend

pytorch · vmoens · Oct 24, 2023 · Oct 6, 2023 · Oct 6, 2023 · Oct 6, 2023
commit 3c41dab26a3ceb6d0fdaef118eed9f0d82780ecf
diff --git a/examples/decision_transformer/utils.py b/examples/decision_transformer/utils.py
@@ -20,15 +20,15 @@
     EnvCreator,
     ExcludeTransform,
     ObservationNorm,
+    ParallelEnv,
     RandomCropTensorDict,
     Reward2GoTransform,
     RewardScaling,
     RewardSum,
-    SerialEnv,
     TargetReturn,
     TensorDictPrimer,
     TransformedEnv,
-    UnsqueezeTransform, ParallelEnv,
+    UnsqueezeTransform,
 )
 from torchrl.envs.libs.dm_control import DMControlEnv
 from torchrl.envs.libs.gym import set_gym_backend

diff --git a/torchrl/data/tensor_specs.py b/torchrl/data/tensor_specs.py
@@ -1782,7 +1782,9 @@ def __init__(
 
         dtype, device = _default_dtype_and_device(dtype, device)
         box = (
-            ContinuousBox(torch.tensor(-np.inf).expand(shape), torch.tensor(np.inf).expand(shape))
+            ContinuousBox(
+                torch.tensor(-np.inf).expand(shape), torch.tensor(np.inf).expand(shape)
+            )
             if shape == _DEFAULT_SHAPE
             else None
         )

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -29,7 +29,6 @@
     BoundedTensorSpec,
     CompositeSpec,
     ContinuousBox,
-    DEVICE_TYPING,
     DiscreteTensorSpec,
     MultiDiscreteTensorSpec,
     MultiOneHotDiscreteTensorSpec,
@@ -642,7 +641,7 @@ def transform(self, transform: Transform):
             prev_transform.__dict__["_container"] = None
         transform.set_container(self)
         transform.eval()
-        self._transform = transform
+        self._transform = transform.to(self.device)
 
     @property
     def device(self) -> bool:
@@ -951,8 +950,9 @@ def __init__(self, *transforms: Transform):
     def to(self, *args, **kwargs):
         # because Module.to(...) does not call to(...) on sub-modules, we have
         # manually call it:
-        for t in self.transforms:
-            t.to(*args, **kwargs)
+        self.transforms = nn.ModuleList(
+            [t.to(*args, **kwargs) for t in self.transforms]
+        )
         return super().to(*args, **kwargs)
 
     def _call(self, tensordict: TensorDictBase) -> TensorDictBase: