How to Use .to(device) in PyTorch for Device Management

In PyTorch, use .to(device) to move tensors or models to a specific device like CPU or GPU. This method ensures your data and computations run on the right hardware by specifying device as 'cpu' or 'cuda'.

Syntax

The .to(device) method moves a tensor or model to the specified device.

• tensor_or_model.to(device): Moves the tensor or model to the device.
• device: A string or torch.device object specifying the target device, e.g., 'cpu' or 'cuda'.

python

tensor_or_model.to(device)

Example

This example shows how to move a tensor and a simple model to GPU if available, otherwise to CPU. It prints the device of the tensor and model parameters.

python

import torch
import torch.nn as nn

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

# Create a tensor
x = torch.tensor([1.0, 2.0, 3.0])
print(f'Original tensor device: {x.device}')

# Move tensor to device
x = x.to(device)
print(f'Tensor device after .to(device): {x.device}')

# Define a simple model
model = nn.Linear(3, 1)
print(f'Original model parameter device: {next(model.parameters()).device}')

# Move model to device
model = model.to(device)
print(f'Model parameter device after .to(device): {next(model.parameters()).device}')

Output

Original tensor device: cpu
Tensor device after .to(device): cuda:0
Original model parameter device: cpu
Model parameter device after .to(device): cuda:0

Common Pitfalls

Common mistakes when using .to(device) include:

• Not moving both model and input tensors to the same device, causing errors during computation.
• Using .to(device) without checking if the device (like GPU) is available.
• Forgetting to assign the result of .to(device) back to the variable, since it does not modify in place.

python

import torch

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

x = torch.tensor([1.0, 2.0])

# Wrong: forgetting to assign back
x.to(device)  # This does NOT move x in place
print(x.device)  # Still cpu

# Right: assign back
x = x.to(device)
print(x.device)  # Now moved to device

Output

cpu
cuda:0

Quick Reference

Tips for using .to(device) effectively:

• Always check device availability with torch.cuda.is_available().
• Move both inputs and models to the same device before training or inference.
• Remember .to(device) returns a new object; assign it back.
• Use torch.device objects for clarity and flexibility.

Key Takeaways

Use .to(device) to move tensors or models to CPU or GPU for correct computation.

Always assign the result of .to(device) back to your variable as it does not change in place.

Check if GPU is available before using 'cuda' device to avoid errors.

Move both your model and input data to the same device to prevent runtime errors.

Use torch.device objects for clear and flexible device management.