amend

pytorch · Feb 15, 2024 · e79ef5e · e79ef5e
1 parent d84f13e
commit e79ef5e
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -1162,7 +1162,6 @@ def step_and_maybe_reset(
         #   and this transform overrides an observation key (eg, CatFrames)
         #   the shape, dtype or device may not necessarily match and writing
         #   the value in-place will fail.
-        assert not self.shared_tensordict_parent.requires_grad
         self.shared_tensordict_parent.update_(
             tensordict, keys_to_update=self._env_input_keys
         )
@@ -1224,7 +1223,6 @@ def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         #   and this transform overrides an observation key (eg, CatFrames)
         #   the shape, dtype or device may not necessarily match and writing
         #   the value in-place will fail.
-        assert not self.shared_tensordict_parent.requires_grad
         self.shared_tensordict_parent.update_(
             tensordict, keys_to_update=list(self._env_input_keys)
         )