amend

pytorch · vmoens · Sep 2, 2023 · Jul 28, 2023 · Jul 28, 2023 · Aug 10, 2023
commit 71d1076791349543bf67334b46963b3059f70abe
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
@@ -4,13 +4,17 @@
 # LICENSE file in the root directory of this source tree.
 
 import argparse
-from mocking_classes import DiscreteActionVecMockEnv
-from tensordict.nn import TensorDictSequential
-from torchrl.modules import MLP, ProbabilisticActor
+
 import pytest
 import torch
+from mocking_classes import DiscreteActionVecMockEnv
 from tensordict import pad, TensorDict, unravel_key_list
-from tensordict.nn import InteractionType, make_functional, TensorDictModule
+from tensordict.nn import (
+    InteractionType,
+    make_functional,
+    TensorDictModule,
+    TensorDictSequential,
+)
 from torch import nn
 from torchrl.data.tensor_specs import (
     BoundedTensorSpec,
@@ -23,6 +27,7 @@
     DecisionTransformerInferenceWrapper,
     DTActor,
     LSTMModule,
+    MLP,
     NormalParamWrapper,
     OnlineDTActor,
     ProbabilisticActor,
@@ -1768,21 +1773,20 @@ def test_multi_consecutive(self, shape):
         )
 
     def test_lstm_parallel_env(self):
-        from torchrl.envs import ParallelEnv, TransformedEnv, InitTracker
+        from torchrl.envs import InitTracker, ParallelEnv, TransformedEnv
+
         # tests that hidden states are carried over with parallel envs
         lstm_module = LSTMModule(
-                input_size=7,
-                hidden_size=12,
-                num_layers=2,
-                in_key="observation",
-                out_key="features",
-            )
+            input_size=7,
+            hidden_size=12,
+            num_layers=2,
+            in_key="observation",
+            out_key="features",
+        )
 
         def create_transformed_env():
             primer = lstm_module.make_tensordict_primer()
-            env = DiscreteActionVecMockEnv(
-                categorical_action_encoding=True
-                )
+            env = DiscreteActionVecMockEnv(categorical_action_encoding=True)
             env = TransformedEnv(env)
             env.append_transform(InitTracker())
             env.append_transform(primer)

diff --git a/torchrl/modules/tensordict_module/rnn.py b/torchrl/modules/tensordict_module/rnn.py
@@ -5,7 +5,7 @@
 from typing import Optional, Tuple
 
 import torch
-from tensordict import unravel_key_list, unravel_key, TensorDictBase
+from tensordict import TensorDictBase, unravel_key_list
 
 from tensordict.nn import TensorDictModuleBase as ModuleBase