diff --git a/tb_plugin/torch_tb_profiler/profiler/memory_parser.py b/tb_plugin/torch_tb_profiler/profiler/memory_parser.py
index 91d2f578f..1b9369547 100644
--- a/tb_plugin/torch_tb_profiler/profiler/memory_parser.py
+++ b/tb_plugin/torch_tb_profiler/profiler/memory_parser.py
@@ -47,6 +47,8 @@ def device_name(self):
             return 'CPU'
         elif self.device_type == DeviceType.CUDA:
             return 'GPU{}'.format(self.device_id)
+        elif self.device_type is not None:
+            return self.device_type.name
         else:
             return None
 
diff --git a/tb_plugin/torch_tb_profiler/profiler/trace.py b/tb_plugin/torch_tb_profiler/profiler/trace.py
index 538a29246..854ee9817 100644
--- a/tb_plugin/torch_tb_profiler/profiler/trace.py
+++ b/tb_plugin/torch_tb_profiler/profiler/trace.py
@@ -17,7 +17,26 @@
 
 class DeviceType(IntEnum):
     CPU = 0
-    CUDA = 1
+    CUDA = 1 # CUDA.
+    MKLDNN = 2 # Reserved for explicit MKLDNN
+    OPENGL = 3 # OpenGL
+    OPENCL = 4 # OpenCL
+    IDEEP = 5 # IDEEP.
+    HIP = 6 # AMD HIP
+    FPGA = 7 # FPGA
+    MAIA = 8 # ONNX Runtime / Microsoft
+    XLA = 9 # XLA / TPU
+    Vulkan = 10 # Vulkan
+    Metal = 11 # Metal
+    XPU = 12 # XPU
+    MPS = 13 # MPS
+    Meta = 14 # Meta (tensors with no data)
+    HPU = 15 # HPU / HABANA
+    VE = 16 # SX-Aurora / NEC
+    Lazy = 17 # Lazy Tensors
+    IPU = 18 # Graphcore IPU
+    MTIA = 19 # Meta training and inference devices
+    PrivateUse1 = 20 # PrivateUse1 device
 
 
 class EventTypes: