[Bugfix] Fix Hunyuan worker device context (#3768)

Signed-off-by: Chenguang ZHENG <645327136@qq.com>
Signed-off-by: Canlin Guo <961750412@qq.com>
Co-authored-by: Canlin Guo <961750412@qq.com>
This commit is contained in:
Chenguang ZhengandCanlin Guo authored and GitHub committed 2026-05-22 02:17:59 +00:00
1 parent 1de15bdf4d
commit 7d8816b46b
1 file changed
+2
+2
View File
@@ -10,5 +10,7 @@ from vllm_omni.platforms import current_omni_platform
def get_local_device() -> torch.device:
"""Return the torch device for the current rank based on detected device type."""
if current_omni_platform.is_initialized():
return current_omni_platform.get_torch_device(current_omni_platform.current_device())
local_rank = int(os.environ.get("LOCAL_RANK", 0))
return current_omni_platform.get_torch_device(local_rank)