From 09cfdbf4cc52a3ca878c8679f73ca5192e71d2f9 Mon Sep 17 00:00:00 2001
From: iefgnoix <isaacwxf23@gmail.com>
Date: Fri, 31 May 2024 17:49:49 -0700
Subject: [PATCH] Reuse DLDeviceType. (#7163)

---
 torch_xla/utils/dlpack.py | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/torch_xla/utils/dlpack.py b/torch_xla/utils/dlpack.py
index c49083e4403..73c7d074bff 100644
--- a/torch_xla/utils/dlpack.py
+++ b/torch_xla/utils/dlpack.py
@@ -1,5 +1,6 @@
 from typing import Any
 import enum
+from torch.utils.dlpack import DLDeviceType
 import torch_xla
 
 
@@ -7,20 +8,6 @@ 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__'):