Modern AI workstations often have multiple GPUs; and you need to know which GPUs are available and how to direct your workload to the right one. In this post we’ll walk through device enumeration with PyTorch, show how to load models onto a specific GPU, and cover practical patterns that apply to any PyTorch project.
On a single-GPU machine, torch.device(“cuda”) is all you need. But the moment a second GPU appears, PyTorch defaults to cuda:0. If your inference server and your training job are all fighting over device 0 while device 1 sits idle, you’re leaving performance on the table.
Explicit device selection lets you:
- Balance memory pressure across GPUs.
- Isolate workloads
- Mix hardware (e.g. load a large UNet on an RTX 5090 while a smaller CLIP model fits on an RTX 4060).
Enumerate available devices
def get_available_torch_devices():
devices =
if torch.cuda.is_available():
for i in range(torch.cuda.device_count()):
devices.append(f"cuda:{i}")
return devices
Setting the device for model loading
Option A: torch.device
The most portable and explicit approach is to pass a torch.device object wherever PyTorch accepts one.
device = torch.device(“cuda:1”)
Load checkpoint weights onto a specific GPU
state_dict = torch.load(“model.pt”, map_location=device, weights_only=True)
Tensors created on the right device
x = torch.randn(1, 3, 512, 512, device=device)
output = model(x)
Option B: torch.cuda.device — scoped context manager
When a library or third-party function doesn’t accept a device argument, you can wrap the call in a context manager to influence the default CUDA device for that block:
with torch.cuda.device(1):
model = ThirdPartyModel.from_pretrained(“weights/”)
The key principle: be explicit about devices at every boundary—file loading, model placement, and tensor creation. This avoids the silent default-to-cuda:0 behavior that can cause out-of-memory errors on one GPU while another sits empty.
** Please note that example assumes a homogeneous GPU configuration