Devices: moving models and data to the GPU

Everything in one operation must share a device

Each tensor lives on exactly one device. An operation whose inputs sit on different devices fails with Expected all tensors to be on the same device. So you move the model once and every batch inside the loop.

The two .to() calls behave differently. model.to(device) moves the module's parameters in place and also returns the module, so model = model.to(device) and plain model.to(device) both work. tensor.to(device) returns a new tensor and leaves the original where it was, so you must reassign: x = x.to(device). When possible, create tensors directly on the device: torch.zeros(4, 2, device=device).

CUDA work is asynchronous: Python queues kernels and moves on. Anything that needs a value back on the CPU, such as .item(), .cpu(), .tolist() or printing a tensor, forces Python to wait for the GPU. A few per epoch are fine; one per batch in a tight loop can noticeably slow training. For the same reason, time GPU code with torch.cuda.synchronize() before reading the clock. If you hit CUDA out of memory, reduce the batch size first.

The path data takes to the GPU

Batches are prepared in CPU host memory, copied into pinned (page-locked) buffers for fast transfer, moved to GPU memory, and processed by GPU kernels, which mixed precision and compilation make faster.

Moving modules and tensors

The module moves in place; the tensor must be reassigned. Creating ones directly on the device avoids an extra copy.

import torch
from torch import nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = nn.Linear(10, 2).to(device)     # module parameters move in place
x = torch.randn(4, 10)
x = x.to(device)                        # tensors return a copy: reassign!
out = model(x)
print(out.device)

ones = torch.ones(4, 2, device=device)  # created directly on the device
print((out + ones).shape)               # torch.Size([4, 2])

Reading the device-mismatch error

When you see Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu, check three suspects: a batch you forgot to move, a tensor created inside forward without device=x.device, and a loss weight tensor (such as CrossEntropyLoss(weight=...)) that stayed on the CPU.

Quick check: A model is on the GPU. You call `x.to("cuda")` without assigning the result, then run `model(x)`. What happens?

• It works, because `.to()` moved `x` in place
• It works, but runs on the CPU
• It fails with a device mismatch, because `x` itself is still on the CPU
• PyTorch moves the model back to the CPU automatically