How to move model to gpu pytorch

PytorchHow-ToBeginner · 3 min read

How to Move Model to GPU in PyTorch: Simple Guide

To move a model to GPU in PyTorch, use model.to('cuda') or model.cuda(). This transfers the model's parameters and buffers to the GPU device, enabling faster computation.

📐

Syntax

Use model.to('cuda') or model.cuda() to move your model to the GPU. The to() method is more flexible and can move the model to any device, while cuda() specifically moves it to the default CUDA device.

model: your PyTorch model instance.
'cuda': string specifying the GPU device.
to(): method to move model to a device.
cuda(): shortcut to move model to GPU.

python

model.to('cuda')
# or
model.cuda()

💻

Example

This example shows how to create a simple model, check if a GPU is available, and move the model to GPU. It also prints the device of the model's parameters to confirm the move.

python

import torch
import torch.nn as nn

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# Check if GPU is available
if torch.cuda.is_available():
    device = torch.device('cuda')
    model = model.to(device)
    print('Model moved to:', next(model.parameters()).device)
else:
    print('GPU not available, model on:', next(model.parameters()).device)

Output

Model moved to: cuda:0

⚠️

Common Pitfalls

Common mistakes when moving models to GPU include:

Not checking if GPU is available before moving the model, which causes errors on machines without GPUs.
Moving the model but forgetting to move input data to the same device, causing device mismatch errors.
Using model.cuda() without specifying device when multiple GPUs exist, which may not select the intended GPU.

Always move both model and data to the same device.

python

import torch

model = torch.nn.Linear(5, 2)

# Wrong: moving model but not input data
model = model.to('cuda')
input_data = torch.randn(1, 5)  # on CPU
# This will cause an error:
# output = model(input_data)  # RuntimeError: Expected all tensors to be on the same device

# Right: move input data to GPU too
input_data = input_data.to('cuda')
output = model(input_data)
print('Output device:', output.device)

Output

Output device: cuda:0

📊

Quick Reference

Action	Code Example	Notes
Move model to GPU	model.to('cuda')	Preferred flexible method
Move model to GPU (shortcut)	model.cuda()	Moves to default GPU device
Check GPU availability	torch.cuda.is_available()	Avoid errors on CPU-only machines
Move input tensor to GPU	input_tensor.to('cuda')	Must match model device
Specify GPU device	model.to('cuda:1')	For multi-GPU setups

✅

Key Takeaways

Always check if GPU is available with torch.cuda.is_available() before moving model.

Use model.to('cuda') to move your model to GPU safely and flexibly.

Move input data to the same device as the model to avoid errors.

model.cuda() is a shortcut but less flexible than model.to('cuda').

For multiple GPUs, specify the device like 'cuda:0' or 'cuda:1'.