amend

pytorch · Feb 2, 2024 · 14c92a7 · 14c92a7
1 parent 4f34621
commit 14c92a7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/torchrl/envs/batched_envs.py b/torchrl/envs/batched_envs.py
@@ -769,7 +769,7 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
         device = self.device
         # select + clone creates 2 tds, but we can create one only
         def select_and_clone(name, tensor):
-            if name in selected_output_keys:
+            if _unravel_key_to_tuple(name) in selected_output_keys:
                 return tensor.clone()
 
         out = self.shared_tensordict_parent.named_apply(
@@ -813,7 +813,6 @@ def _step(
         def select_and_clone(name, tensor):
             if _unravel_key_to_tuple(name) in self._selected_step_keys:
                 return tensor.clone()
-            print(f"{name} not in {self._selected_step_keys}")
 
         out = next_td.named_apply(
             select_and_clone, nested_keys=True