Add TorchFix linter to CI (meta-pytorch#569)

Summary:
Add TorchFix (https://github.com/pytorch/test-infra/tree/main/tools/torchfix) flake8 plugin.

Pull Request resolved: meta-pytorch#569

Test Plan: Verify that buck2 test torchtnt/tests/utils:test_module_summary passes.

Reviewed By: galrotem

Differential Revision: D50438385

Pulled By: kit1980

fbshipit-source-id: 68d7d07d2ff3e4374ea5e5e91896720b85671ba6

Author: kit1980

.flake8

@@ -1,6 +1,6 @@
 [flake8]
 # Suggested config from pytorch that we can adopt
-select = B,C,E,F,P,T4,W,B9
+select = B,C,E,F,P,T4,W,B9,TOR0,TOR1,TOR2
 max-line-length = 120
 # C408 ignored because we like the dict keyword argument syntax
 # E501 is not flexible enough, we're using B950 instead

.pre-commit-config.yaml

@@ -27,6 +27,8 @@ repos:
     -   id: flake8
         args:
         - --config=.flake8
+        additional_dependencies:
+          - torchfix==0.1.1
 
 -   repo: https://github.com/omnilib/ufmt
     rev: v1.3.0

tests/utils/test_module_summary.py

@@ -9,13 +9,13 @@
 from typing import Any, Dict, Optional, Tuple
 
 import torch
-import torchvision.models as models
 from torchtnt.utils.module_summary import (
     _get_human_readable_count,
     get_module_summary,
     ModuleSummary,
     prune_module_summary,
 )
+from torchvision import models
 
 
 def get_summary_and_prune(
@@ -155,7 +155,9 @@ def test_lazy_tensor_flops(self) -> None:
 
     def test_resnet_max_depth(self) -> None:
         """Test the behavior of max_depth on a layered model like ResNet"""
-        pretrained_model = models.resnet.resnet18(pretrained=True)
+        pretrained_model = models.resnet.resnet18(
+            weights=models.ResNet18_Weights.IMAGENET1K_V1  # pyre-ignore[16]
+        )
 
         # max_depth = None
         ms1 = get_module_summary(pretrained_model)
@@ -211,7 +213,9 @@ def test_module_summary_layer_print(self) -> None:
         self._test_module_summary_text(summary_table, str(ms1))
 
     def test_alexnet_print(self) -> None:
-        pretrained_model = models.alexnet(pretrained=True)
+        pretrained_model = models.alexnet(
+            weights=models.AlexNet_Weights.IMAGENET1K_V1  # pyre-ignore[16]
+        )
         ms1 = get_summary_and_prune(pretrained_model, max_depth=1)
         ms2 = get_summary_and_prune(pretrained_model, max_depth=2)
         ms3 = get_summary_and_prune(pretrained_model, max_depth=3)
@@ -236,7 +240,9 @@ def test_alexnet_print(self) -> None:
         self.assertEqual(str(ms3), str(ms4))
 
     def test_alexnet_with_input_tensor(self) -> None:
-        pretrained_model = models.alexnet(pretrained=True)
+        pretrained_model = models.alexnet(
+            weights=models.AlexNet_Weights.IMAGENET1K_V1  # pyre-ignore[16]
+        )
         inp = torch.randn(1, 3, 224, 224)
         ms1 = get_summary_and_prune(pretrained_model, max_depth=1, module_args=(inp,))
         ms2 = get_summary_and_prune(pretrained_model, max_depth=2, module_args=(inp,))
@@ -344,7 +350,9 @@ def forward(self, x, y, offset=1):
         self.assertEqual(ms_classifier.out_size, [1, 1, 224, 224])
 
     def test_forward_elapsed_time(self) -> None:
-        pretrained_model = models.alexnet(pretrained=True)
+        pretrained_model = models.alexnet(
+            weights=models.AlexNet_Weights.IMAGENET1K_V1  # pyre-ignore[16]
+        )
         inp = torch.randn(1, 3, 224, 224)
         ms1 = get_summary_and_prune(pretrained_model, module_args=(inp,), max_depth=4)
         stack = [ms1] + [