Correcting index test cas

Author: apbose

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

@@ -9,6 +9,7 @@
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
 from torch_tensorrt.dynamo.conversion.converter_utils import (
     broadcastable,
+    cast_trt_tensor,
     get_positive_dim,
     get_trt_tensor,
     to_numpy,
@@ -81,16 +82,16 @@ def gather(
     name: str,
     input: TRTTensor,
     dim: int,
-    index: Sequence[Union[TRTTensor, np.ndarray, torch.Tensor]],
+    index: Union[TRTTensor, np.ndarray, torch.Tensor],
     sparse_grad: bool = False,
 ) -> TRTTensor:
-    indices_tensor = []
-
-    for i, ind in enumerate(index):
-        indices_tensor.append(
-            get_trt_tensor(ctx, ind, name + f"_parameter_to_fp32_tensor_{i}")
-        )
-    gather_layer = ctx.net.add_gather(input, indices_tensor, dim)
+    if not isinstance(index, TRTTensor):
+        index = get_trt_tensor(ctx, index, name + f"_parameter_to_fp32_tensor")
+    # This is for the case where torch.ops.aten.gather requires torch.int64
+    # However TRTInterpreter complains that torch.int64 is not a supported type
+    # So the below cast does not help
+    # index = cast_trt_tensor(ctx, input, trt.int32, name, target, source_ir)
+    gather_layer = ctx.net.add_gather(input, index, dim)
     set_layer_name(gather_layer, target, name + "_gather", source_ir)
     return gather_layer.get_output(0)
 

tests/py/dynamo/conversion/test_gather_aten.py

@@ -2,26 +2,25 @@
 
 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
 
+from .harness import DispatchTestCase
+
 
 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))
                 super().__init__()
 
             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
 
-        index0 = torch.randint(0, 1, (1, 1), dtype=torch.int32)
-        indices = [None, index0]
+        # index0 = torch.randint(0, 1, (1, 1), dtype=torch.int32)
+        index0 = torch.randint(0, 1, (1, 1))
         input = [torch.randn(2, 2), index0]
         self.run_test(
             TestModule(),