lint

pytorch · Jan 30, 2024 · abe7436 · abe7436
1 parent f66c73b
commit abe7436
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/test/mocking_classes.py b/test/mocking_classes.py
@@ -229,7 +229,8 @@ def _step(self, tensordict):
                 "terminated": done.clone(),
                 "observation": n.clone(),
             },
-            batch_size=[], device=self.device,
+            batch_size=[],
+            device=self.device,
         )
 
     def _reset(self, tensordict: TensorDictBase = None, **kwargs) -> TensorDictBase:
@@ -241,7 +242,9 @@ def _reset(self, tensordict: TensorDictBase = None, **kwargs) -> TensorDictBase:
         done = self.counter >= self.max_val
         done = torch.tensor([done], dtype=torch.bool, device=self.device)
         return TensorDict(
-            {"done": done, "terminated": done.clone(), "observation": n}, [], device=self.device,
+            {"done": done, "terminated": done.clone(), "observation": n},
+            [],
+            device=self.device,
         )
 
     def rand_step(self, tensordict: Optional[TensorDictBase] = None) -> TensorDictBase: