11 lines
297 B
Python
11 lines
297 B
Python
|
# SPDX-License-Identifier: Apache-2.0
|
||
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||
|
|
|
||
|
|
import os
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def get_local_device() -> torch.device:
|
||
|
|
"""Return the torch device for the current rank."""
|
||
|
|
return torch.device(f"cuda:{os.environ.get('LOCAL_RANK', 0)}")
|