From 28193fd98548d9f8373b1c2e492565874ae61c76 Mon Sep 17 00:00:00 2001
From: Xinyu Yang <yangxinyu.302@bytedance.com>
Date: Thu, 16 May 2024 15:33:23 +0800
Subject: [PATCH] [Stablehlo]index type use i64 (#3354)

---
 lib/Conversion/TorchToStablehlo/GatherScatter.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lib/Conversion/TorchToStablehlo/GatherScatter.cpp b/lib/Conversion/TorchToStablehlo/GatherScatter.cpp
index 00c022cc1..5854b1b7d 100644
--- a/lib/Conversion/TorchToStablehlo/GatherScatter.cpp
+++ b/lib/Conversion/TorchToStablehlo/GatherScatter.cpp
@@ -247,8 +247,7 @@ FailureOr<Value> broadcastAndConcatIndices(Operation *op,
   concatShape.push_back(indexTensors.size());
 
   SmallVector<Value> broadcastedIndices;
-  Type indexElemTy =
-      cast<RankedTensorType>(indexTensors[0].getType()).getElementType();
+  Type indexElemTy = rewriter.getI64Type();
   RankedTensorType bcastIndexType =
       RankedTensorType::get(indicesShape, indexElemTy);
   for (auto indexTensor : indexTensors) {