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.

1. 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}")

2. 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.

3. 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()

4. 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.

5. 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.

6. 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.

7. Remove the smoke-test script after the project code uses the selector.

   $ rm select_device_demo.py