Fix fx2trt CI test trigger condition (#71014)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/71014

Replace test trigger with test_config matching.

Test Plan:
CI
https://github.com/pytorch/pytorch/runs/4746717568?check_suite_focus=true

Reviewed By: janeyx99

Differential Revision: D33480971

fbshipit-source-id: 9513e464753343a7ae47fcfaf48119f34bae94c5

.jenkins/pytorch/test.sh

@@ -568,7 +568,7 @@ elif [[ "${BUILD_ENVIRONMENT}" == *distributed* || "${JOB_BASE_NAME}" == *distri
   test_rpc
 elif [[ "${TEST_CONFIG}" = docs_test ]]; then
   test_docs_test
-elif [[ "${BUILD_ENVIRONMENT}" == *linux-xenial-cuda11.3-py3.6-gcc7* ]]; then
+elif [[ "${TEST_CONFIG}" == fx2trt ]]; then
   test_torch_fx2trt
 else
   install_torchvision

test/fx2trt/converters/acc_op/test_adaptive_avgpool.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from parameterized import parameterized
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestAdaptiveAvgPoolConverter(AccTestCase):
@@ -48,3 +49,6 @@ class TestAdaptiveAvgPoolConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.adaptive_avg_pool2d}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_avgpool.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from parameterized import parameterized, param
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestAvgPoolConverter(AccTestCase):
@@ -45,3 +46,6 @@ class TestAvgPoolConverter(AccTestCase):
 
         inputs = [torch.randn(1, 3, 224, 224)]
         self.run_test(TestModule(), inputs, expected_ops={acc_ops.avg_pool2d})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_batchnorm.py

@@ -3,6 +3,7 @@
 import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestBatchNormConverter(AccTestCase):
@@ -38,3 +39,6 @@ class TestBatchNormConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.batch_norm}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_binary_ops.py

@@ -7,6 +7,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 elementwise_ops = [
     ((lambda x, y: x + y), acc_ops.add),
@@ -101,3 +102,6 @@ class TestBinaryOpConverters(AccTestCase):
         ]
 
         self.run_test_with_dynamic_shape(Op(), input_specs, expected_ops={expected_op})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_cat.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestCatConverter(AccTestCase):
@@ -34,3 +35,6 @@ class TestCatConverter(AccTestCase):
             ),
         ]
         self.run_test_with_dynamic_shape(Cat(), input_specs, expected_ops={acc_ops.cat})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_chunk.py

@@ -3,6 +3,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from parameterized import parameterized
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestChunkConverter(AccTestCase):
@@ -24,3 +25,6 @@ class TestChunkConverter(AccTestCase):
             inputs,
             expected_ops={acc_ops.chunk},
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_clamp.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized, param
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestClampConverter(AccTestCase):
@@ -27,3 +28,6 @@ class TestClampConverter(AccTestCase):
 
         inputs = [torch.randn(3, 4)]
         self.run_test(TestModule(), inputs, expected_ops={acc_ops.clamp})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_convolution.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized, param
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestConvolutionConverter(AccTestCase):
@@ -59,3 +60,6 @@ class TestConvolutionConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.conv2d}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_dequantize.py

@@ -7,6 +7,7 @@ import torch.fx
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 @unittest.skip(
@@ -44,3 +45,6 @@ class TestDequantizeConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.dequantize}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_elu.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestELUConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestELUConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.elu}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_flatten.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestFlattenConverter(AccTestCase):
@@ -65,3 +66,6 @@ class TestFlattenConverter(AccTestCase):
             input_specs,
             expected_ops={acc_ops.flatten},
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_gelu.py

@@ -2,6 +2,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestGELU(AccTestCase):
@@ -33,3 +34,6 @@ class TestGELU(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.gelu}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_getitem.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestGetitemConverter(AccTestCase):
@@ -52,3 +53,6 @@ class TestGetitemConverter(AccTestCase):
 
         inputs = [torch.randn(2, 10, 10, 10)]
         self.run_test(Getitem(idx), inputs, expected_ops={acc_ops.getitem})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_hardtanh.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestHardtanhConverter(AccTestCase):
@@ -14,3 +15,6 @@ class TestHardtanhConverter(AccTestCase):
 
         inputs = [torch.randn(2, 10, 10, 10)]
         self.run_test(Hardtanh(), inputs, expected_ops={acc_ops.hardtanh})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_layernorm.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized, param
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestLayerNormConverter(AccTestCase):
@@ -59,3 +60,6 @@ class TestLayerNormConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             LayerNorm(normalized_shape), input_specs, expected_ops={acc_ops.layer_norm}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_leaky_relu.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestLeakyReLUConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestLeakyReLUConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.leaky_relu}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_linear.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestLinearConverter(AccTestCase):
@@ -50,3 +51,6 @@ class TestLinearConverter(AccTestCase):
             input_specs,
             expected_ops={acc_ops.linear},
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_matmul.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMatMulConverter(AccTestCase):
@@ -85,3 +86,6 @@ class TestMatMulConverter(AccTestCase):
             expected_ops={acc_ops.matmul},
             test_implicit_batch_dim=test_implicit_batch_dim,
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_max.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMaxConverter(AccTestCase):
@@ -77,3 +78,6 @@ class TestMaxConverter(AccTestCase):
 
         inputs = [torch.randn(3, 4), torch.randn(3, 4)]
         self.run_test(MaxMethod(), inputs, expected_ops={acc_ops.maximum})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_maximum.py

@@ -3,6 +3,7 @@
 import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMaximumConverter(AccTestCase):
@@ -29,3 +30,6 @@ class TestMaximumMethodConverter(AccTestCase):
             torch.randn(3, 4),
         ]
         self.run_test(Maximum(), inputs, expected_ops={acc_ops.maximum})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_maxpool.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized, param
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMaxPoolConverter(AccTestCase):
@@ -58,3 +59,6 @@ class TestMaxPoolConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.max_pool2d}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_min.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMinConverter(AccTestCase):
@@ -77,3 +78,6 @@ class TestMinConverter(AccTestCase):
 
         inputs = [torch.randn(3, 4), torch.randn(3, 4)]
         self.run_test(MinMethod(), inputs, expected_ops={acc_ops.minimum})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_minimum.py

@@ -3,6 +3,7 @@
 import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestMinimumConverter(AccTestCase):
@@ -29,3 +30,6 @@ class TestMinimumMethodConverter(AccTestCase):
             torch.randn(3, 4),
         ]
         self.run_test(Minimum(), inputs, expected_ops={acc_ops.minimum})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_narrow.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestNarrowConverter(AccTestCase):
@@ -26,3 +27,6 @@ class TestNarrowConverter(AccTestCase):
             expected_ops={acc_ops.slice_tensor},
             test_explicit_batch_dim=False,
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_pad.py

@@ -7,6 +7,7 @@ from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import param, parameterized
 import unittest
 import tensorrt as trt
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestPadConverter(AccTestCase):
@@ -67,3 +68,6 @@ class TestPadConverter(AccTestCase):
             inputs,
             expected_ops={acc_ops.pad},
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_permute.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestPermuteConverter(AccTestCase):
@@ -51,3 +52,6 @@ class TestPermuteConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             Permute(), input_specs, expected_ops={acc_ops.permute}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_quantize_per_tensor.py

@@ -7,6 +7,7 @@ import torch.fx
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 @unittest.skip(
@@ -42,3 +43,6 @@ class TestQuantizePerTensorConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.quantize_per_tensor}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_reduce_ops.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 reduce_ops = [(torch.sum, acc_ops.sum), (torch.mean, acc_ops.mean)]
 
@@ -63,3 +64,6 @@ class TestReduceConverter(AccTestCase):
             expected_ops={expected_acc_op},
             test_implicit_batch_dim=False,
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_relu.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestReLUConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestReLUConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.relu}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_reshape.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestReshapeConverter(AccTestCase):
@@ -50,3 +51,6 @@ class TestReshapeConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(target_shape), input_specs, expected_ops={acc_ops.reshape}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_selu.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSeLUConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestSeLUConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.selu}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_sigmoid.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSigmoid(AccTestCase):
@@ -14,3 +15,6 @@ class TestSigmoid(AccTestCase):
 
         inputs = [torch.randn(1, 2, 3)]
         self.run_test(Sigmoid(), inputs, expected_ops={acc_ops.sigmoid})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_size.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSizeConverter(AccTestCase):
@@ -32,3 +33,6 @@ class TestSizeConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             Size(), input_specs, expected_ops={acc_ops.size}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_softmax.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSoftmaxConverter(AccTestCase):
@@ -64,3 +65,6 @@ class TestSoftmaxConverter(AccTestCase):
             expect_error=AssertionError,
             test_explicit_batch_dim=False,
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_softsign.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSoftsignConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestSoftsignConverter(AccTestCase):
         self.run_test_with_dynamic_shape(
             TestModule(), input_specs, expected_ops={acc_ops.softsign}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_split.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSplitConverter(AccTestCase):
@@ -30,3 +31,6 @@ class TestSplitConverter(AccTestCase):
             },
             test_explicit_batch_dim=False,
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_squeeze.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestSqueeze(AccTestCase):
@@ -30,3 +31,6 @@ class TestSqueeze(AccTestCase):
         self.run_test_with_dynamic_shape(
             Squeeze(), input_specs, expected_ops={acc_ops.squeeze}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_tanh.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestTanh(AccTestCase):
@@ -30,3 +31,6 @@ class TestTanh(AccTestCase):
         self.run_test_with_dynamic_shape(
             Tanh(), input_specs, expected_ops={acc_ops.tanh}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_tile.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestTile(AccTestCase):
@@ -68,3 +69,6 @@ class TestTile(AccTestCase):
         self.run_test_with_dynamic_shape(
             Tile(dims), input_specs, expected_ops={acc_ops.tile}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_topk.py

@@ -5,6 +5,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestTopKConverter(AccTestCase):
@@ -41,3 +42,6 @@ class TestTopKConverter(AccTestCase):
             expected_ops={acc_ops.topk},
             test_implicit_batch_dim=(dim != 0),
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_unary_ops.py

@@ -7,6 +7,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 unary_ops = [
     (torch.sin, acc_ops.sin),
@@ -43,3 +44,6 @@ class TestUnaryOpConverters(AccTestCase):
         m = TestModule(orig_op)
         inputs = [torch.randn(2, 2, 3)]
         self.run_test(m, inputs, expected_ops={expected_op})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/acc_op/test_unsqueeze.py

@@ -6,6 +6,7 @@ import torch.fx.experimental.fx_acc.acc_ops as acc_ops
 import torch.nn as nn
 from torch.testing._internal.common_fx2trt import AccTestCase, InputTensorSpec
 from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestUnsqueeze(AccTestCase):
@@ -52,3 +53,6 @@ class TestUnsqueeze(AccTestCase):
         self.run_test_with_dynamic_shape(
             Unsqueeze(dim), input_specs, expected_ops={acc_ops.unsqueeze}
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/vanilla/test_add.py

@@ -5,6 +5,7 @@ import operator
 import torch
 import torch.fx
 from torch.testing._internal.common_fx2trt import VanillaTestCase
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestAddConverter(VanillaTestCase):
@@ -21,3 +22,6 @@ class TestAddConverter(VanillaTestCase):
 
         inputs = [torch.randn(1, 1)]
         self.run_test(add, inputs, expected_ops={torch.add})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/converters/vanilla/test_convolution.py

@@ -4,6 +4,7 @@ import torch
 import torch.fx
 from torch.testing._internal.common_fx2trt import VanillaTestCase
 from parameterized import parameterized, param
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestConvolutionConverter(VanillaTestCase):
@@ -39,3 +40,6 @@ class TestConvolutionConverter(VanillaTestCase):
 
         inputs = [torch.randn(1, 3, 224, 224)]
         self.run_test(TestModule(), inputs, expected_ops={torch.nn.modules.conv.Conv2d})
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/core/test_import_fx2trt.py

@@ -1,13 +1,16 @@
 #!/usr/bin/env python3
 # (c) Facebook, Inc. and its affiliates. Confidential and proprietary.
 
-import unittest as ut
 
 # Test that this import should not trigger any error when run
 # in non-GPU hosts, or in any build mode.
 import torch.fx.experimental.fx2trt.lower as fxl  # noqa: F401
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
-class MainTests(ut.TestCase):
+class MainTests(TestCase):
     def test_1(self):
         pass
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/core/test_input_tensor_spec.py

@@ -1,13 +1,13 @@
 from typing import List, Optional
-import unittest
 
 import torch
 from torch.fx.experimental.fx2trt import (
     InputTensorSpec,
 )
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
-class TestTRTModule(unittest.TestCase):
+class TestTRTModule(TestCase):
     def _validate_spec(
         self,
         spec: InputTensorSpec,
@@ -44,3 +44,6 @@ class TestTRTModule(unittest.TestCase):
             for batch_size, shape in zip(batch_size_range, spec.shape_ranges[0]):
                 self.assertEqual(batch_size, shape[0])
                 self.assertSequenceEqual(tensor.shape[1:], shape[1:])
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/core/test_trt_module.py

@@ -1,5 +1,3 @@
-import unittest
-
 import torch
 import torch.fx
 import torch.fx.experimental.fx_acc.acc_tracer as acc_tracer
@@ -8,9 +6,11 @@ from torch.fx.experimental.fx2trt import (
     InputTensorSpec,
     TRTModule,
 )
+from torch.testing._internal.common_utils import TestCase, run_tests
+import os
 
 
-class TestTRTModule(unittest.TestCase):
+class TestTRTModule(TestCase):
     def test_save_and_load_trt_module(self):
         class TestModule(torch.nn.Module):
             def forward(self, x):
@@ -29,6 +29,7 @@ class TestTRTModule(unittest.TestCase):
         torch.testing.assert_allclose(
             reload_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
         )
+        os.remove(f"{os.getcwd()}/trt.pt")
 
     def test_save_and_load_state_dict(self):
         class TestModule(torch.nn.Module):
@@ -50,3 +51,6 @@ class TestTRTModule(unittest.TestCase):
         torch.testing.assert_allclose(
             new_trt_mod(inputs[0].cuda()).cpu(), ref_output, rtol=1e-04, atol=1e-04
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/passes/test_fuse_permute_linear_trt.py

@@ -5,6 +5,7 @@ from torch.fx.experimental.fx2trt.passes.fuse_pass import (
     fuse_permute_linear,
     trt_transposed_linear,
 )
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestFusePermuteLinear(AccTestCase):
@@ -72,3 +73,6 @@ class TestFusePermuteLinear(AccTestCase):
             {trt_transposed_linear},
             apply_passes=[fuse_permute_linear],
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/passes/test_fuse_permute_matmul_trt.py

@@ -6,6 +6,7 @@ from torch.fx.experimental.fx2trt.passes.fuse_pass import (
     fuse_permute_matmul,
     trt_transposed_matmul,
 )
+from torch.testing._internal.common_utils import run_tests
 
 
 def tranpose_last_two_dims(x):
@@ -118,3 +119,6 @@ class TestFusePermuteMatmul(AccTestCase):
             {trt_transposed_matmul},
             apply_passes=[fuse_permute_matmul],
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/passes/test_fuse_unsqueeze_cat_sum_trt.py

@@ -4,6 +4,7 @@ from torch.testing._internal.common_fx2trt import AccTestCase
 from torch.fx.experimental.fx2trt.passes.fuse_pass import (
     fuse_unsqueeze_cat_sum,
 )
+from torch.testing._internal.common_utils import run_tests
 
 
 class TestFuseUnsqueezeCatSum(AccTestCase):
@@ -33,3 +34,6 @@ class TestFuseUnsqueezeCatSum(AccTestCase):
             {acc_ops.unsqueeze, acc_ops.cat, acc_ops.sum},
             apply_passes=[fuse_unsqueeze_cat_sum],
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/passes/test_multi_fuse_trt.py

@@ -8,6 +8,7 @@ from torch.fx.experimental.fx2trt.passes.fuse_pass import (
     fuse_permute_matmul,
     trt_transposed_matmul,
 )
+from torch.testing._internal.common_utils import run_tests
 
 
 def permute021(x):
@@ -57,3 +58,6 @@ class TestMultiFuse(AccTestCase):
             {acc_ops.permute},
             apply_passes=[fuse_permute_matmul, fuse_permute_linear],
         )
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/passes/test_remove_duplicate_output_args.py

@@ -1,15 +1,15 @@
 import logging
-import unittest as ut
 
 import torch.fx as fx
 import torch.fx.experimental.fx2trt.passes.remove_duplicate_output_args as dedup
 import torch.nn as nn
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class TestFx2TrtPasses(ut.TestCase):
+class TestFx2TrtPasses(TestCase):
     def test_remove_duplicate_output_args(self):
         class Sub(nn.Module):
             def forward(self, x):
@@ -65,3 +65,6 @@ graph():
         assert (
             ttop_a_graph_expected == ttop_a_graph_actual
         ), f"Unexpected ttop.a graph: {ttop_a_graph_actual}"
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/trt_lower/test_fx2trt_lower.py

@@ -1,17 +1,17 @@
 import logging
-import unittest as ut
 # @manual=//caffe2:torch_fx2trt
 from torch.fx.experimental.fx2trt.lower import Lowerer
 from torch.fx.experimental.fx2trt import LowerSetting
 import torch
 import torch.fx as fx
 import torch.nn as nn
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
 logger = logging.getLogger(__name__)
 
 
-class Fx2trtLowerTests(ut.TestCase):
+class Fx2trtLowerTests(TestCase):
     def test_fx2trt_lower(self):
         mod = _Mod()
         mod_traced = fx.symbolic_trace(mod)
@@ -24,3 +24,6 @@ class Fx2trtLowerTests(ut.TestCase):
 class _Mod(nn.Module):
     def forward(self, x):
         return (x, 2 * x)
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/trt_lower/trt_operator_supported_test.py

@@ -1,15 +1,15 @@
 # Owner(s): ["oncall: fx"]
 
-import unittest
 import torch.fx.experimental.fx_acc.acc_ops  # noqa: F401
 import torch
 import torch.fx
 from torch.fx.experimental.fx2trt.tools.trt_splitter import create_trt_operator_support
 import torch.nn as nn
 from torch.fx.experimental.fx_acc import acc_ops, acc_tracer
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
-class TestTRTOperatorSupport(unittest.TestCase):
+class TestTRTOperatorSupport(TestCase):
     def test_supported_node_target(self):
         class TestModule(nn.Module):
             def __init__(self):
@@ -76,3 +76,6 @@ class TestTRTOperatorSupport(unittest.TestCase):
         for node in traced_mod.graph.nodes:
             if node.target == acc_ops.quantize_per_tensor:
                 self.assertTrue(op_support.is_node_supported(mod, node))
+
+if __name__ == '__main__':
+    run_tests()

test/fx2trt/trt_lower/trt_splitter_test.py

@@ -1,7 +1,6 @@
 # Owner(s): ["oncall: fx"]
 
 import operator
-import unittest
 
 import torch  # isort:skip
 import torch.fx  # isort:skip
@@ -12,6 +11,7 @@ import torch.fx.passes.shape_prop as shape_prop
 from torch.fx.experimental.fx2trt.tools.trt_splitter import TRTSplitter
 from torch.fx.passes import splitter_base
 from torch.fx.experimental.fx_acc import acc_tracer
+from torch.testing._internal.common_utils import TestCase, run_tests
 
 
 ERROR_MSG_NO_ACC_MODULE = "FX split failed: Did not find any ACC submodule!"
@@ -61,7 +61,7 @@ def find_call_targets(module: torch.fx.GraphModule):
 # We test both FxNetSplitOnly and FxNetSplitter here, since they share most
 # functionalities. The only difference is that FxNetSplitOnly does not implement
 # split_preview() related functions, while FxNetSplitter does.
-class TestSplit(unittest.TestCase):
+class TestSplit(TestCase):
     def test_demo(self):
         """
           ==> b ==>
@@ -643,7 +643,7 @@ class TestSplit(unittest.TestCase):
         test_splitter(splitter)
 
 
-class TestSplitComplexGraph(unittest.TestCase):
+class TestSplitComplexGraph(TestCase):
     """
            a ======
         //   \\     \\
@@ -753,7 +753,7 @@ class TestSplitComplexGraph(unittest.TestCase):
         test_splitter(splitter)
 
 
-class TestSplitNonTensorEdges(unittest.TestCase):
+class TestSplitNonTensorEdges(TestCase):
     """
            a (relu)
         //   \\
@@ -963,7 +963,7 @@ class TestSplitNonTensorEdges(unittest.TestCase):
         test_splitter(splitter)
 
 
-class TestAccNodesFinder(unittest.TestCase):
+class TestAccNodesFinder(TestCase):
     def test_acc_nodes_finder_1(self):
         """
         y ------------->
@@ -1007,7 +1007,7 @@ class TestAccNodesFinder(unittest.TestCase):
         self.assertEqual(set(), acc_nodes, "Shouldn't have ACC nodes")
 
 
-class TestAccFusionsFinder(unittest.TestCase):
+class TestAccFusionsFinder(TestCase):
     """
               x
              / \\
@@ -1147,3 +1147,6 @@ class TestAccFusionsFinder(unittest.TestCase):
 
 def op_support_with_support_dict(support_dict: dict) -> op_support.OperatorSupportBase:
     return op_support.OperatorSupport(support_dict)
+
+if __name__ == '__main__':
+    run_tests()

test/run_test.py

@@ -893,7 +893,7 @@ def get_selected_tests(options):
     # Only run fx2trt test with specified option argument
     if options.fx2trt_tests:
         selected_tests = list(
-            filter(lambda test_name: "fx2trt" in test_name, selected_tests)
+            filter(lambda test_name: test_name in FX2TRT_TESTS, selected_tests)
         )
     else:
         options.exclude.extend(FX2TRT_TESTS)

torch/testing/_internal/common_fx2trt.py

@@ -38,6 +38,7 @@ def fetch_attr(mod, target):
 @unittest.skipIf(not torch.cuda.is_available(), "Skip because CUDA is not available")
 class TRTTestCase(TestCase):
     def setUp(self):
+        super().setUp()
         torch.manual_seed(3)
 
     def run_test(self, mod, inputs, expected_ops, unexpected_ops, interpreter, rtol, atol):