lint

Author: chenhao388

test/test_nn.py

@@ -59,7 +59,6 @@
 from tensordict.tensorclass import TensorClass
 from _utils_internal import is_npu_available
 
-
 from torch import distributions, nn
 from torch.distributions import Categorical, Normal
 from torch.utils._pytree import tree_map
@@ -83,7 +82,6 @@
 except ImportError:
     from tensordict.utils import Buffer
 
-
 IS_FB = os.getenv("PYTORCH_TEST_FBCODE")
 
 # Capture all warnings
@@ -113,6 +111,7 @@
         pytest.mark.filterwarnings("ignore:inplace"),
     )
 
+
 def get_device():
     device = torch.device("cpu")
     if torch.cuda.is_available():
@@ -123,6 +122,7 @@ def get_device():
         device = torch.device("mps:0")
     return device
 
+
 class TestInteractionType:
     def test_base(self):
         with set_interaction_type("DETERMINISTIC"):
@@ -2161,7 +2161,7 @@ def test_module_buffer():
         module.cuda()
         assert module.td.device.type == "cuda"
     elif is_npu_available():
-        module.npu()
+        module = module.to("npu:0")
         assert module.td.device.type == "npu"
 
 
@@ -2186,7 +2186,6 @@ def test_module_buffer():
 )
 def test_to_context(original_device, new_device, tc):
     if tc:
-
         class MyTC(TensorClass):
             x: torch.Tensor
             y: torch.Tensor | None = None

test/test_tensorclass.py

@@ -739,7 +739,6 @@ def test_disallowed_attributes(self):
             AttributeError,
             match="Attribute name reshape can't be used with @tensorclass",
         ):
-
             @tensorclass
             class MyInvalidClass:
                 x: torch.Tensor
@@ -1101,7 +1100,6 @@ class MyDataParent:
     @pytest.mark.parametrize("list_to_stack", [True, False])
     def test_indexing(self, list_to_stack):
         with set_list_to_stack(list_to_stack):
-
             @tensorclass
             class MyDataNested:
                 X: torch.Tensor
@@ -1438,8 +1436,8 @@ class MyDataNested(TensorClass):
         assert (
             repeated.X
             == X.repeat_interleave(
-                torch.tensor([2, 3, 4, 5], device=data.device), dim=1
-            )
+            torch.tensor([2, 3, 4, 5], device=data.device), dim=1
+        )
         ).all()
 
     def test_reshape(self):
@@ -2890,23 +2888,20 @@ class FuncAutoCast:
 class TestShadow:
     def test_no_shadow(self):
         with pytest.raises(AttributeError):
-
             @tensorclass
             class MyClass:
                 x: str
                 y: int
                 batch_size: Any
 
         with pytest.raises(AttributeError):
-
             @tensorclass
             class MyClass:  # noqa: F811
                 x: str
                 y: int
                 names: Any
 
         with pytest.raises(AttributeError):
-
             @tensorclass
             class MyClass:  # noqa: F811
                 x: str
@@ -3104,7 +3099,7 @@ class MyClass:
         _ = c / 1
         _ = 1 / c
 
-        _ = c**1
+        _ = c ** 1
         # not implemented
         # 1 ** c
 
@@ -3304,15 +3299,13 @@ class TensorOnly:
             c: torch.Tensor | None = None
 
         with pytest.raises(TypeError, match="tensor_only"):
-
             @tensorclass(tensor_only=True, nocast=True)
             class TensorOnlyNocast:
                 a: torch.Tensor
                 b: torch.Tensor
                 c: torch.Tensor | None = None
 
         with pytest.raises(TypeError, match="tensor_only"):
-
             @tensorclass(tensor_only=True, autocast=True)
             class TensorOnlyAutocast:
                 a: torch.Tensor
@@ -3337,7 +3330,6 @@ class TensorOnly(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
-
             class TensorOnlyAny(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: Any
@@ -3347,7 +3339,6 @@ class TensorOnlyAny(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
-
             class TensorOnlyStr(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: torch.Tensor | str
@@ -3357,7 +3348,6 @@ class TensorOnlyStr(TensorClass["tensor_only"]):
             TypeError,
             match="tensor_only requires types to be Tensor, Tensor-subtrypes or None",
         ):
-
             class TensorOnlyStrUnion(TensorClass["tensor_only"]):
                 a: torch.Tensor
                 b: torch.Tensor