fix api mistakes

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

Author: xuechendi

vllm_omni/platforms/cuda/platform.py

@@ -118,8 +118,8 @@ def get_device_name(cls, device_id: int = 0) -> str:
 
     @classmethod
     def reset_peak_memory_stats(cls, device: torch.device | None = None) -> None:
-        torch.xpu.reset_peak_memory_stats(device)
+        torch.cuda.reset_peak_memory_stats(device)
 
     @classmethod
     def current_device_index(cls) -> int:
-        return torch.xpu.current_device()
+        return torch.cuda.current_device()

vllm_omni/platforms/npu/platform.py

@@ -84,11 +84,11 @@ def get_free_memory(cls, device: torch.device | None = None) -> int:
 
     @classmethod
     def reset_peak_memory_stats(cls, device: torch.device | None = None) -> None:
-        torch.xpu.reset_peak_memory_stats(device)
+        torch.npu.reset_peak_memory_stats(device)
 
     @classmethod
     def current_device_index(cls) -> int:
-        return torch.xpu.current_device()
+        return torch.npu.current_device()
 
     @classmethod
     def get_device_total_memory(cls, device_id: int = 0) -> int:

vllm_omni/platforms/rocm/platform.py

@@ -102,8 +102,8 @@ def get_free_memory(cls, device: torch.device | None = None) -> int:
 
     @classmethod
     def reset_peak_memory_stats(cls, device: torch.device | None = None) -> None:
-        torch.xpu.reset_peak_memory_stats(device)
+        torch.cuda.reset_peak_memory_stats(device)
 
     @classmethod
     def current_device_index(cls) -> int:
-        return torch.xpu.current_device()
+        return torch.cuda.current_device()