Correct sparse arg in aten::gather

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -204,7 +204,7 @@ def aten_ops_gather(
         input=args[0],
         dim=args[1],
         index=args[2],
-        sparse_grad=args_bounds_check(args, 4, False),
+        sparse_grad=args_bounds_check(args, 3, False),
     )
 
 

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

@@ -82,6 +82,7 @@ def gather(
     input: TRTTensor,
     dim: int,
     index: Sequence[Union[TRTTensor, np.ndarray, torch.Tensor]],
+    sparse_grad: bool = False,
 ) -> TRTTensor:
     gather_layer = ctx.net.add_gather(input, index, dim)
     set_layer_name(gather_layer, target, name + "_gather", source_ir)

tests/py/dynamo/conversion/test_gather_aten.py

@@ -2,25 +2,27 @@
 
 import torch
 import torch.nn as nn
-from harness import DispatchTestCase
+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 TestGatherConverter(DispatchTestCase):
+    def test_gather_zero_two_dim(self):
         class TestModule(nn.Module):
             def __init__(self):
-                self.index0 = torch.randint(0, 1, (1, 1))
+                # 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]
+            def forward(self, x, indices):
+                # 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)]
+        index0 = torch.randint(0, 1, (1, 1), dtype=torch.int32)
+        indices = [None, index0]
+        input = [torch.randn(2, 2), index0]
         self.run_test(
             TestModule(),
             input,