Modifications to matmul and select tests

Author: apbose

py/torch_tensorrt/fx/test/converters/aten_op/test_matmul_aten.py

@@ -11,7 +11,7 @@ def test_matmul(self):
         class TestModule(torch.nn.Module):
             def forward(self, x, y):
                 return torch.matmul(x, y)
-        inputOne = torch.randn(1, 32)
+        inputOne = torch.randn(3, 32)
         inputTwo = torch.randn(32, 3)
         inputs = [inputOne, inputTwo]
         self.run_test(

py/torch_tensorrt/fx/test/converters/aten_op/test_select_aten.py

@@ -7,19 +7,25 @@
 from torch_tensorrt.fx.tools.common_fx2trt import DispatchTestCase, InputTensorSpec
 
 class TestSelectConverter(DispatchTestCase):
-    def test_select(self):
+    @parameterized.expand(
+        [
+            ("select_dim_index", 2, 1),
+        ]
+    )
+    def test_select(self, dim_test, index_test):
         class TestModule(torch.nn.Module):
-            def forward(self, input, dim, index):
-                return torch.select(input, dim, index)
+            def __init__(self, dim, index):
+                super().__init__()
+                self.dim = dim
+                self.index = index
+            def forward(self, input):
+                return torch.select(input, self.dim, self.index)
         input = [torch.randn(1, 3, 32)]
-        dim = 2
-        index = 1
-        inputs = (input, dim, index)
         self.run_test(
-            TestModule(), input, expected_ops={torch.ops.aten.select.Tensor}, test_explicit_precision=True,
+            TestModule(dim_test, index_test), input, expected_ops={torch.ops.aten.select}, test_explicit_precision=True,
         )
 
-    def test_select_with_dynamic_shape(self, x, y):
+    def test_select_with_dynamic_shape(self, dim_test, index_test):
         class TestModule(torch.nn.Module):
             def forward(self, input, dim, index):
                 return torch.select(input, dim, index)
@@ -31,9 +37,9 @@ def forward(self, input, dim, index):
                 shape_ranges=[((1, 3, 3), (3, 3, 3), (32, 32, 32))],
             ),
         ]
-        dim = 2
-        index = 1
-        inputs_spec = (input_spec, dim, index)
         self.run_test_with_dynamic_shape(
-            TestModule(), inputs_spec, expected_ops={torch.ops.aten.select.Tensor}
-        )
+            TestModule(dim_test, index_test), input_spec, expected_ops={torch.ops.aten.select}
+        )
+
+if __name__ == "__main__":
+    run_tests()