PyTorch Tensor on GPU

Training a deep learning model on a CPU is like processing 10,000 math problems one at a time. A GPU does those same 10,000 problems in parallel. Moving your tensors to a GPU is one of the most impactful steps you can take to speed up training.

CPU vs GPU — The Key Difference

CPU                          GPU
──────────────────────────────────────────────────
4 to 16 cores (powerful)     Thousands of cores (simpler)
Good at sequential tasks     Good at parallel tasks
Runs all general code        Specialized for math arrays
Slow for large matrix math   Fast for large matrix math

Neural network training is dominated by matrix multiplications — exactly the kind of math that GPUs excel at. Moving your model and tensors to the GPU can make training 10x to 100x faster depending on the task.

Checking GPU Availability

import torch

print(torch.cuda.is_available())       # True if NVIDIA GPU is ready
print(torch.cuda.device_count())       # number of GPUs
print(torch.cuda.get_device_name(0))   # GPU name, e.g. "NVIDIA GeForce RTX 3080"

Setting the Device

The standard pattern is to set a device variable and use it throughout your code. This way, the same code works on CPU-only machines and GPU machines without changes.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)   # either "cuda" or "cpu"

Moving Tensors to GPU

import torch

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

t = torch.tensor([1.0, 2.0, 3.0])
print(t.device)      # cpu

t_gpu = t.to(device)
print(t_gpu.device)  # cuda:0  (if GPU available)

Diagram: CPU → GPU → CPU Flow

[Data on Disk]
      │
      ▼ load with DataLoader
[Tensor on CPU]
      │
      ▼ .to(device)
[Tensor on GPU]      ← training happens here (fast)
      │
      ▼ .cpu() or .to('cpu')
[Tensor on CPU]      ← needed for NumPy, plotting, saving
      │
      ▼
[Save results / plot]

Moving a Model to GPU

Moving the model to the GPU works the same way as moving a tensor. After this, all computations inside the model run on the GPU.

import torch
import torch.nn as nn

model = nn.Linear(10, 1)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)
print(next(model.parameters()).device)   # cuda:0

Running a Full GPU Training Step

import torch
import torch.nn as nn

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

model = nn.Linear(10, 1).to(device)
optimizer = torch.optim.Adam(model.parameters())
loss_fn = nn.MSELoss()

# Data must also be on the same device as the model
X = torch.rand(32, 10).to(device)   # batch of 32 samples, 10 features
y = torch.rand(32, 1).to(device)    # target values

pred = model(X)
loss = loss_fn(pred, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()

print("Loss:", loss.item())

The Most Common GPU Error

If you send your model to the GPU but forget to send your data, you get a RuntimeError about tensors being on different devices. The fix is simple: always move both model and data to the same device.

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

Fix:
  model → .to(device)
  data  → .to(device)
  ✓ both on same device

Moving Tensors Back to CPU

NumPy does not support GPU tensors. Before converting a PyTorch GPU tensor to a NumPy array, you must move it back to the CPU first.

t_gpu = torch.tensor([1.0, 2.0, 3.0]).to("cuda")

# This fails: t_gpu.numpy()

# Fix:
t_cpu = t_gpu.cpu()
arr = t_cpu.detach().numpy()
print(arr)   # [1. 2. 3.]

Apple Silicon (MPS) Support

If you use a Mac with an Apple M1, M2, or M3 chip, PyTorch supports GPU acceleration through Metal Performance Shaders (MPS). Use mps instead of cuda.

device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

Multi-GPU Setup

If you have multiple GPUs, you can target a specific one by its index.

device = torch.device("cuda:0")   # first GPU
device = torch.device("cuda:1")   # second GPU

For training across all available GPUs simultaneously, PyTorch provides nn.DataParallel and DistributedDataParallel — covered in the advanced topics section of this course.

Google Colab Free GPU

If you don't have a GPU on your machine, Google Colab gives you free access to NVIDIA GPUs. In Colab, go to Runtime → Change runtime type → Hardware accelerator → GPU. Your code runs exactly the same way — just with cuda available.

Summary

GPUs accelerate neural network training by running thousands of math operations in parallel. Check GPU availability with torch.cuda.is_available(). Set a device variable and move both your model and your data to it using .to(device). Always keep the model and data on the same device to avoid runtime errors. Before converting tensors to NumPy, move them back to CPU with .cpu(). Apple Mac users can use MPS for GPU acceleration on Apple Silicon.