small fix: Index validator enable int64

- Repair test case

Author: gs-olive

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -397,7 +397,7 @@ def index_dtype_validator(node: Node) -> bool:
     for ind in index:
         if ind is not None:
             val = ind.meta.get("val")
-            if val is not None and val.dtype != torch.int32:
+            if val is not None and val.dtype not in (torch.int32, torch.int64):
                 return False
     return True
 

tests/py/dynamo/conversion/test_index_aten.py

@@ -1,10 +1,8 @@
-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
+
+from .harness import DispatchTestCase
 
 
 class TestIndexConverter(DispatchTestCase):
@@ -15,7 +13,6 @@ def __init__(self):
                 super().__init__()
 
             def forward(self, x):
-                index0 = torch.randint(0, 1, (1, 1))
                 indices = [None, self.index0]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out
@@ -158,8 +155,6 @@ def __init__(self):
                 super().__init__()
 
             def forward(self, x):
-                index0 = torch.tensor([0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7])
-                index1 = index0.unsqueeze(0).T.long()
                 indices = [None, None, self.index0, self.index1]
                 out = torch.ops.aten.index.Tensor(x, indices)
                 return out