Linting fix and adding test case

Author: apbose

py/torch_tensorrt/dynamo/conversion/impl/select.py

@@ -253,7 +253,9 @@ def index(
                     dim_tensor_list[adv_indx_indices[i]],
                 )
 
-        gather_out = gather(ctx, target, source_ir, name, flatten_tensor, 0, cum_adv_index)
+        gather_out = gather(
+            ctx, target, source_ir, name, flatten_tensor, 0, cum_adv_index
+        )
         _LOGGER.debug(f"The shape after cumultative gather is {gather_out.shape}")
         _LOGGER.debug(f"The shape for cumulative adv index is {cum_adv_index}")
 

tests/py/dynamo/conversion/test_gather_aten.py

@@ -0,0 +1,30 @@
+import operator
+
+import torch
+import torch.nn as nn
+from harness import DispatchTestCase
+from torch.testing._internal.common_utils import run_tests
+from torch_tensorrt import Input
+
+
+class TestIndexConverter(DispatchTestCase):
+    def test_index_zero_two_dim(self):
+        class TestModule(nn.Module):
+            def __init__(self):
+                self.index0 = torch.randint(0, 1, (1, 1))
+                super().__init__()
+
+            def forward(self, x):
+                index0 = torch.randint(0, 1, (1, 1))
+                indices = [None, self.index0]
+                out = torch.ops.aten.gather.default(x, 0, indices)
+                return out
+
+        input = [torch.randn(2, 2)]
+        self.run_test(
+            TestModule(),
+            input,
+        )
+
+if __name__ == "__main__":
+    run_tests()