DataType check

deepjavalibrary · May 11, 2022 · 706ec85 · 706ec85
1 parent ea84044
commit 706ec85
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -347,12 +347,18 @@ public static void set(PtNDArray self, ByteBuffer data) {
     }
 
     public static PtNDArray gather(PtNDArray ndArray, PtNDArray index, long dim) {
+        if (index.getDataType() != DataType.INT64) {
+            index = index.toType(DataType.INT64, true);
+        }
         return new PtNDArray(
                 ndArray.getManager(),
                 PyTorchLibrary.LIB.torchGather(ndArray.getHandle(), index.getHandle(), dim, false));
     }
 
     public static PtNDArray take(PtNDArray ndArray, PtNDArray index) {
+        if (index.getDataType() != DataType.INT64) {
+            index = index.toType(DataType.INT64, true);
+        }
         return new PtNDArray(
                 ndArray.getManager(),
                 PyTorchLibrary.LIB.torchTake(ndArray.getHandle(), index.getHandle()));