amend

pytorch · Oct 2, 2023 · 559d157 · 559d157
1 parent 1bb5745
commit 559d157
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -2080,7 +2080,9 @@ def __init__(
     ):
         if in_keys is None:
             in_keys = IMAGE_KEYS
-        super(GrayScale, self).__init__(in_keys=in_keys, out_keys=out_keys)
+        if out_keys is None:
+            out_keys = copy(in_keys)
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
 
     def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
         observation = F.rgb_to_grayscale(observation)