How to beginner · 3 min read

How to use GPU with PyTorch

Q: How to use GPU with PyTorch

Use torch.cuda.is_available() to check GPU availability and move your tensors and models to GPU with .to('cuda') or .cuda(). This enables PyTorch to leverage GPU acceleration for faster computation.

Quick answer

Use torch.cuda.is_available() to check GPU availability and move your tensors and models to GPU with .to('cuda') or .cuda(). This enables PyTorch to leverage GPU acceleration for faster computation.

PREREQUISITES

Python 3.8+
PyTorch installed with CUDA support (e.g. pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu117)
A CUDA-capable GPU and appropriate NVIDIA drivers installed

Setup

Install PyTorch with CUDA support matching your GPU and CUDA version. Verify your NVIDIA drivers and CUDA toolkit are installed. Use the official PyTorch installation selector at pytorch.org to get the correct install command.

bash

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu117

Step by step

This example checks for GPU availability, moves a tensor and a model to GPU, and performs a simple tensor operation on the GPU.

python

import torch

# Check if GPU is available
if torch.cuda.is_available():
    device = torch.device('cuda')
    print(f'Using GPU: {torch.cuda.get_device_name(0)}')
else:
    device = torch.device('cpu')
    print('GPU not available, using CPU')

# Create a tensor and move it to the GPU
x = torch.randn(3, 3)
x = x.to(device)
print('Tensor device:', x.device)

# Define a simple model and move it to GPU
model = torch.nn.Linear(3, 1)
model = model.to(device)

# Forward pass with tensor on GPU
output = model(x)
print('Output:', output)
print('Output device:', output.device)

output

Using GPU: NVIDIA GeForce RTX 3080
Tensor device: cuda:0
Output: tensor([[...]], device='cuda:0', grad_fn=<AddmmBackward0>)
Output device: cuda:0

Common variations

Use .cuda() as shorthand for .to('cuda').
For multi-GPU setups, specify device IDs like torch.device('cuda:1').
Use torch.nn.DataParallel or torch.nn.parallel.DistributedDataParallel for multi-GPU training.
Check GPU memory usage with torch.cuda.memory_allocated() and torch.cuda.memory_reserved().

Troubleshooting

If torch.cuda.is_available() returns False, verify your CUDA drivers and PyTorch CUDA version match.
Restart your machine after installing NVIDIA drivers.
Ensure your GPU supports the CUDA version installed.
Use nvidia-smi in terminal to check GPU status.
If you get RuntimeError: CUDA out of memory, reduce batch size or clear cache with torch.cuda.empty_cache().

✅

Key Takeaways

Always check GPU availability with torch.cuda.is_available() before using GPU.
Move both your model and tensors to the GPU device using .to('cuda') or .cuda().
Use torch.device objects to write device-agnostic code that works on CPU or GPU.
For multi-GPU, specify device IDs and consider parallel wrappers like DataParallel.
Troubleshoot CUDA errors by verifying driver versions, GPU compatibility, and memory usage.

Verified 2026-04

Verify ↗