Moving tensors between CPU and GPU

Setup#

pip install torch

This lesson runs correctly with or without a GPU. Outputs below were captured with torch 2.5.1 on a machine with an NVIDIA card; the noted lines differ on CPU-only machines.

The portable pattern#

import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print("using:", device)

x = torch.ones(3)             # tensors are born on the CPU
print(x.device)
x = x.to(device)              # .to() COPIES it to the device
print(x.device)

y = torch.ones(3, device=device)   # or create it there directly
z = x + y                     # both operands on one device: allowed
print(z)

back = z.cpu().numpy()        # NumPy only understands CPU memory
print(back)

using: cuda
cpu
cuda:0
tensor([2., 2., 2.], device='cuda:0')
[2. 2. 2.]

On a machine without a GPU, the same file prints using: cpu, every device reads cpu, and no line crashes. That is the point of the first line — write it once at the top and the code runs everywhere.

The walkthrough#

x = x.to(device) — note the assignment. .to() returns a new tensor on the target device; it does not move the original. x.to(device) alone, without the =, is the most common single-line GPU bug in existence: it computes a copy and throws it away.

cuda:0 means "the first GPU" — the :0 is an index, visible when a machine has several.

Mixing devices fails loudly:

a = torch.ones(3)
b = torch.ones(3, device="cuda")
print(a + b)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

You will meet this message many times. It nearly always means one tensor missed the ferry — often the labels while the inputs made it, or a freshly created tensor inside the loop.

z.cpu().numpy() — the return trip. NumPy, matplotlib, pandas and printing large results all need CPU data. Calling .numpy() on a CUDA tensor raises an error telling you to call .cpu() first.

Models move the same way: model.to(device) moves every parameter inside (in place, unlike tensors — one of PyTorch's stranger asymmetries). Then each batch moves in the loop: inputs = inputs.to(device).

Common mistakes#

Moving inside the hot loop unnecessarily. Shipping the model every batch, or bouncing results back per step for printing, drowns the GPU's speed in ferry rides. Move the model once; accumulate metrics on-device and bring one number back per epoch.

Timing without waiting. GPU work runs asynchronously — Python queues the work and races ahead. Timing with time.time() right after an operation measures the queueing, not the work. Call torch.cuda.synchronize() before reading the clock.

Assuming the GPU is always faster. For tiny tensors, transfer plus launch overhead exceeds the compute. A GPU earns its keep on big tensors and big models; a 10-element regression is faster at home on the CPU.

Hardcoding "cuda". Code with a literal "cuda" inside crashes on every machine without one — including your reader's. The device = ... if ... else line at the top costs nothing and makes the file portable.

Loading GPU-saved checkpoints on a CPU box. torch.load(path) tries to restore tensors to their original device and fails on machines without one. Pass map_location="cpu".

Try it yourself#

If you have a GPU: create two (4000, 4000) random tensors and matrix-multiply them on CPU, then on GPU, timing both — with a synchronize() before each clock read. Then repeat at (50, 50) and watch the ranking flip. No GPU? Predict both outcomes, and keep the portable pattern.

What to learn next#

• Seeds and reproducible runs — the other global state your experiments depend on.
• Installing PyTorch with the right CUDA build — if cuda.is_available() disagreed with your hardware.
• PyTorch basics — the training loop this pattern slots into.