XLA downcast for S64 and U64 for Neuron (#8347)

pytorch · Nov 7, 2024 · 03f07e2 · 03f07e2
1 parent 6454b42
commit 03f07e2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torch_xla/csrc/dtype.cpp b/torch_xla/csrc/dtype.cpp
@@ -143,9 +143,11 @@ xla::PrimitiveType MaybeDowncastToXlaDeviceType(
       return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
                                         : xla::PrimitiveType::S16;
     case xla::PrimitiveType::S64:
-      return xla::PrimitiveType::S64;
+      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
+                                        : xla::PrimitiveType::S64;
     case xla::PrimitiveType::U64:
-      return xla::PrimitiveType::U64;
+      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::U32
+                                        : xla::PrimitiveType::U64;
     case xla::PrimitiveType::C128:
       return xla::PrimitiveType::C128;
     default: