How to select a device in PyTorch

Machine learning workloads often move between developer laptops, hosted notebooks, and accelerator servers. A portable PyTorch entry point should choose from the backends the active build can initialize instead of assuming every host exposes the same hardware.

The torch.accelerator.current_accelerator() API represents the accelerator compiled into the active PyTorch build, including CUDA or ROCm, MPS, XPU, and other registered backends. Its runtime availability check returns no device when that accelerator cannot initialize, allowing CPU to remain the fallback.

A torch.device object names an allocation target but does not move existing data by itself. Model parameters, input tensors, and output tensors must resolve to the same device type before the forward pass can succeed.

1. Create device_select.py with a runtime accelerator check and CPU fallback.

   device_select.py

   import torch
    
    
   accelerator = torch.accelerator.current_accelerator(check_available=True)
   device = accelerator or torch.device("cpu")

   The check_available=True argument includes a runtime hardware and driver check. This check can initialize accelerator state before fork-based multiprocessing starts.

2. Append a model allocated on the selected device to device_select.py.

   model = torch.nn.Linear(3, 2).to(device)

3. Append an input tensor allocated on the selected device to device_select.py.

   inputs = torch.ones(4, 3, device=device)

4. Append a forward pass to device_select.py.

   with torch.no_grad():
       outputs = model(inputs)

5. Append a fail-capable device alignment guard to device_select.py.

   device_types = {
       next(model.parameters()).device.type,
       inputs.device.type,
       outputs.device.type,
   }
   if device_types != {device.type}:
       raise RuntimeError(f"Device mismatch: {sorted(device_types)}")

6. Append the selected model, input, and output device report to device_select.py.

   print(f"selected_device={device.type}")
   print(f"model_device={next(model.parameters()).device.type}")
   print(f"input_device={inputs.device.type}")
   print(f"output_device={outputs.device.type}")

7. Run the completed device selection script from its project directory.

   $ python device_select.py
   selected_device=cpu
   model_device=cpu
   input_device=cpu
   output_device=cpu

   A CPU-only build prints cpu on every line. An accelerator build reports its registered device type, such as cuda, mps, or xpu.