Rewording of docstring

pytorch · Sep 19, 2024 · 3f620a8 · 3f620a8
1 parent 4622937
commit 3f620a8
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -9046,7 +9046,6 @@ def test_single_trans_env_check(self, create_copy, compose):
         if compose:
             t = Compose(t)
         env = TransformedEnv(ContinuousActionVecMockEnv(), t)
-
         check_env_specs(env)
 
     def test_serial_trans_env_check(self, create_copy):

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -6634,15 +6634,16 @@ def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
 
 
 class RenameTransform(Transform):
-    """A transform to rename entries in the output tensordict.
+    """A transform to rename entries in the output tensordict (or input tensordict
+    via the inverse keys).
 
     Args:
-        in_keys (sequence of NestedKey): the entries to rename
+        in_keys (sequence of NestedKey): the entries to rename.
         out_keys (sequence of NestedKey): the name of the entries after renaming.
-        in_keys_inv (sequence of NestedKey, optional): the entries to rename before
-            passing the input tensordict to :meth:`EnvBase._step`.
-        out_keys_inv (sequence of NestedKey, optional): the names of the renamed
-            entries passed to :meth:`EnvBase._step`.
+        in_keys_inv (sequence of NestedKey, optional): the entries to rename
+            in the input tensordict, which will be passed to :meth:`EnvBase._step`.
+        out_keys_inv (sequence of NestedKey, optional): the names of the entries
+            in the input tensordict after renaming.
         create_copy (bool, optional): if ``True``, the entries will be copied
             with a different name rather than being renamed. This allows for
             renaming immutable entries such as ``"reward"`` and ``"done"``.