Skip to content

GitLab

Unverified Commit 36f6aede authored by Andreas Karatzas's avatar Andreas Karatzas Committed by GitHub
Browse files

[Misc] Optimized check to encapsulate both CUDA and ROCm platforms (#34549) 


Signed-off-by: default avatarAndreas Karatzas <akaratza@amd.com>
parent 9704a5c3
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
vllm/utils/torch_utils.py
View file @ 36f6aede
...@@ -683,7 +683,7 @@ def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tens ...@@ -683,7 +683,7 @@ def get_accelerator_view_from_cpu_tensor(cpu_tensor: torch.Tensor) -> torch.Tens
if current_platform.is_xpu(): if current_platform.is_xpu():
assert cpu_tensor.is_pinned(), "CPU tensor must be pinned" assert cpu_tensor.is_pinned(), "CPU tensor must be pinned"
return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor) return torch.ops._C.get_xpu_view_from_cpu_tensor(cpu_tensor)
elif current_platform.is_cuda() or current_platform.is_rocm(): elif current_platform.is_cuda_alike():
return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor) return torch.ops._C.get_cuda_view_from_cpu_tensor(cpu_tensor)
else: else:
raise ValueError( raise ValueError(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment