Revert S64 and U64 type change for Neuron device due to hardcoded dty…

…pe checks in torch-xla (#8375)
pytorch · Nov 13, 2024 · bc227f7 · bc227f7
1 parent baa14f3
commit bc227f7
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/torch_xla/csrc/dtype.cpp b/torch_xla/csrc/dtype.cpp
@@ -143,11 +143,9 @@ xla::PrimitiveType MaybeDowncastToXlaDeviceType(
       return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
                                         : xla::PrimitiveType::S16;
     case xla::PrimitiveType::S64:
-      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::S32
-                                        : xla::PrimitiveType::S64;
+      return xla::PrimitiveType::S64;
     case xla::PrimitiveType::U64:
-      return CheckNeuronDevice(hw_type) ? xla::PrimitiveType::U32
-                                        : xla::PrimitiveType::U64;
+      return xla::PrimitiveType::U64;
     case xla::PrimitiveType::C128:
       return xla::PrimitiveType::C128;
     default: