Reuse DLDeviceType. (pytorch#7163)

AlibabaPAI · Oct 11, 2024 · 97e28f3 · 97e28f3
1 parent c792f4c
commit 97e28f3
Showing 1 changed file with 1 addition and 14 deletions.
diff --git a/torch_xla/utils/dlpack.py b/torch_xla/utils/dlpack.py
@@ -1,26 +1,13 @@
 from typing import Any
 import enum
+from torch.utils.dlpack import DLDeviceType
 import torch_xla
 
 
 def to_dlpack(xla_tensor: Any):
   return torch_xla._XLAC._to_dlpack(xla_tensor)
 
 
-class DLDeviceType(enum.IntEnum):
-  # Enums as in DLPack specification (aten/src/ATen/dlpack.h)
-  kDLCPU = 1,
-  kDLGPU = 2,
-  kDLCPUPinned = 3,
-  kDLOpenCL = 4,
-  kDLVulkan = 7,
-  kDLMetal = 8,
-  kDLVPI = 9,
-  kDLROCM = 10,
-  kDLExtDev = 12,
-  kDLOneAPI = 14,
-
-
 def from_dlpack(ext_tensor: Any):
   if hasattr(ext_tensor, '__dlpack_device__') and hasattr(
       ext_tensor, '__dlpack__'):