Upcast int tensor indices

Author: honno

torch_np/_ndarray.py

@@ -410,9 +410,19 @@ def clip(self, min, max, out=None):
     )
 
     ### indexing ###
+    @staticmethod
+    def _upcast_int_indices(index):
+        if isinstance(index, torch.Tensor):
+            if index.dtype in [torch.int8, torch.int16, torch.int32]:
+                return index.type(torch.int64)
+        elif isinstance(index, tuple):
+            return tuple(ndarray._upcast_int_indices(i) for i in index)
+        return index
+
     def __getitem__(self, index):
-        t_index = _helpers.ndarrays_to_tensors(index)
-        return ndarray._from_tensor_and_base(self._tensor.__getitem__(t_index), self)
+        index = _helpers.ndarrays_to_tensors(index)
+        index = ndarray._upcast_int_indices(index)
+        return ndarray._from_tensor_and_base(self._tensor.__getitem__(index), self)
 
     def __setitem__(self, index, value):
         value = asarray(value).get()