How to select a device in PyTorch

PyTorch code often runs across laptops, CPU servers, NVIDIA GPUs, Apple Silicon Macs, and notebook runtimes with different accelerator support. Selecting a device at startup keeps the same training or inference code from hard-coding CUDA on machines that only have CPU or MPS available.

torch.device names the allocation target for tensors and modules. torch.cuda.is_available() and torch.backends.mps.is_available() let the program choose an accelerator only when the installed build and hardware can use it, then fall back to cpu.

A selected device only helps after model parameters and input tensors are moved to it. Keep the selection function close to the training or inference entry point, and check at least one parameter, input, and output device during the first run.

Steps to select a PyTorch device:

Create a smoke script that chooses CUDA, MPS, or CPU and runs one model call on that device.

select_device_demo.py

import torch
from torch import nn
 
 
def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")
 
 
device = select_device()
model = nn.Linear(3, 2).to(device)
inputs = torch.ones(4, 3, device=device)
 
with torch.no_grad():
    outputs = model(inputs)
 
print(f"torch_version={torch.__version__}")
print(f"cuda_available={torch.cuda.is_available()}")
print(f"mps_available={torch.backends.mps.is_available()}")
print(f"selected_device={device}")
print(f"model_device={next(model.parameters()).device}")
print(f"input_device={inputs.device}")
print(f"output_device={outputs.device}")

Run the device-selection smoke script.
```
$ python select_device_demo.py
torch_version=2.12.1+cpu
cuda_available=False
mps_available=False
selected_device=cpu
model_device=cpu
input_device=cpu
output_device=cpu
```
The exact torch_version line changes with the active environment. The selected_device, model_device, input_device, and output_device lines should name the same device.

Add the device selector to the project entry point.

def select_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")
 
 
device = select_device()

Move the model to the selected device before creating the optimizer.
```
model = build_model().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
```
Module.to(device) moves model parameters and buffers in place. Keeping the optimizer after the move avoids stale parameter references in accelerator-specific code paths.
Move every input batch to the same device before the forward pass.
```
for inputs, targets in train_loader:
    inputs = inputs.to(device)
    targets = targets.to(device)
 
    outputs = model(inputs)
    loss = loss_fn(outputs, targets)
```
Tensor.to(device) returns a tensor on the requested device. Assign the returned tensor before passing it into the model or loss function.
Print the first run's devices when wiring the selector into existing code.
```
print(f"model_device={next(model.parameters()).device}")
print(f"input_device={inputs.device}")
print(f"output_device={outputs.device}")
```
Matching device values confirm that the model parameters, batch tensors, and output tensors are on the same backend.
Remove the smoke-test script after the project code uses the selector.
```
$ rm select_device_demo.py
```