fix nested out keys

pytorch · Jan 4, 2024 · a5cdeb4 · a5cdeb4
1 parent 2716806
commit a5cdeb4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -6681,7 +6681,7 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
 
         # Update out TensorDict with the burnt-in data.
         for out_key in self.out_keys:
-            if out_key not in td_out.keys():
+            if out_key not in td_out.keys(include_nested=True):
                 td_out.set(
                     out_key,
                     torch.zeros(