[Performance] Reduce key accessing in transforms (#1590)

Signed-off-by: Matteo Bettini <matbet@meta.com>
pytorch · Oct 2, 2023 · 3785609 · 3785609
1 parent 59d29b8
commit 3785609
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -260,8 +260,9 @@ def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
 
         """
         for in_key, out_key in zip(self.in_keys, self.out_keys):
-            if in_key in tensordict.keys(include_nested=True):
-                observation = self._apply_transform(tensordict.get(in_key))
+            value = tensordict.get(in_key, default=None)
+            if value is not None:
+                observation = self._apply_transform(value)
                 tensordict.set(
                     out_key,
                     observation,