Fix GPUMemMonitor callback when using xpu device (#4141)

fix callback when using xpu
openvinotoolkit · Dec 4, 2024 · 57a4f18 · 57a4f18
1 parent ca6b88b
commit 57a4f18
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/otx/algo/callbacks/gpu_mem_monitor.py b/src/otx/algo/callbacks/gpu_mem_monitor.py
@@ -28,7 +28,7 @@ def _get_and_log_device_stats(
             batch_size (int): batch size.
         """
         device = trainer.strategy.root_device
-        if device.type == "cpu":
+        if device.type in ["cpu", "xpu"]:
             return
 
         device_stats = trainer.accelerator.get_device_stats(device)