ruff

Author: jerryzh168

tutorials/developer_api_guide/my_dtype_tensor_subclass.py

@@ -394,9 +394,17 @@ def _(func, types, args, kwargs):
 ########
 # Test #
 ########
-def test():
+def main():
     from torchao.utils import benchmark_model
 
+    class M(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.linear = torch.nn.Linear(1024, 128)
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.linear(x)
+
     m = M()
     example_inputs = (100 * torch.randn(512, 1024),)
     NUM_WARMUPS = 10
@@ -431,4 +439,4 @@ def test():
     print("after quantization and compile:", benchmark_model(m, NUM_RUNS, example_inputs))
 
 if __name__ == "__main__":
-    test()
+    main()

tutorials/developer_api_guide/my_trainable_tensor_subclass.py

@@ -61,7 +61,7 @@ def from_float(
         return _ToMyTrainableDTypeTensor.apply(input_float, layout_type)
 
 class _ToMyTrainableDTypeTensor(torch.autograd.Function):
-    """ 
+    """
     Differentiable constructor for `MyTrainableDTypeTensor`.
     """
 
@@ -163,8 +163,8 @@ def _(func, types, args, kwargs):
 ########
 
 class M(torch.nn.Module):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self) -> None:
+        super().__init__()
         self.linear = torch.nn.Linear(512, 1024, bias=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:

tutorials/developer_api_guide/tensor_parallel.py

@@ -145,7 +145,7 @@ def rowwise_shard(m: torch.nn.Module, mesh: DeviceMesh) -> torch.nn.Module:
 ########
 # Test #
 ########
-if __name__ == "__main__":
+def main():
     # To make sure different ranks create the same module
     torch.manual_seed(5)
 
@@ -192,3 +192,6 @@ def rowwise_shard(m: torch.nn.Module, mesh: DeviceMesh) -> torch.nn.Module:
     print("torch.compile works!")
 
     dist.destroy_process_group()
+
+if __name__ == "__main__":
+    main()