set dtype for xp_take_along_axis

Signed-off-by: Jinzhe Zeng <[email protected]>
njzjz · Oct 23, 2024 · 993b41d · 993b41d
1 parent 6074d33
commit 993b41d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py
@@ -61,7 +61,7 @@ def xp_take_along_axis(arr, indices, axis):
     else:
         indices = xp.reshape(indices, (0, 0))
 
-    offset = (xp.arange(indices.shape[0]) * m)[:, xp.newaxis]
+    offset = (xp.arange(indices.shape[0], dtype=indices.type) * m)[:, xp.newaxis]
     indices = xp.reshape(offset + indices, (-1,))
 
     out = xp.take(arr, indices)