Moving Tensors Across Devices(CPU/GPU)
Anonymous contributor
Published Jan 28, 2025
Contribute to Docs
PyTorch provides simple methods to transfer tensors between CPU and GPU devices, allowing for flexible computation strategies. This is particularly useful when certain operations are more efficient on specific devices, or when working with models that require GPU acceleration. The transfer maintains the tensor’s data and structure while changing only its location in memory.
Syntax
tensor.to(device) # Transfers tensor to the specified device
tensor.cuda() # Transfers tensor to GPU
tensor.cpu() # Transfers tensor to CPU
device
: The destination device, which can be:- A string, e.g.,
'cuda'
,'cpu'
. - A
torch.device
object, e.g.,torch.device('cuda')
ortorch.device('cpu')
.
- A string, e.g.,
Example
The following example demonstrates moving a tensor between CPU and GPU, showing device location at each step, while checking for GPU availability to ensure code works on all systems:
import torch# Create a tensor on CPUcpu_tensor = torch.randn(2, 3)print("Original tensor (CPU):", cpu_tensor.device)# Move to GPU if availableif torch.cuda.is_available():gpu_tensor = cpu_tensor.to('cuda')print("Moved to GPU:", gpu_tensor.device)# Move back to CPUback_to_cpu = gpu_tensor.cpu()print("Back to CPU:", back_to_cpu.device)
The output of the above code will be if the gpu is not available:
Original tensor (CPU): cpu
If gpu is available then:
Original tensor (CPU): cpuMoved to GPU: cuda:0Back to CPU: cpu
Here, cuda:0
indicates the tensor is on the first GPU device.
All contributors
- Anonymous contributor
Contribute to Docs
- Learn more about how to get involved.
- Edit this page on GitHub to fix an error or make an improvement.
- Submit feedback to let us know how we can improve Docs.
Learn PyTorch on Codecademy
- Career path
Computer Science
Looking for an introduction to the theory behind programming? Master Python while learning data structures, algorithms, and more!Includes 6 CoursesWith Professional CertificationBeginner Friendly75 hours - Free course
Intro to PyTorch and Neural Networks
Learn how to use PyTorch to build, train, and test artificial neural networks in this course.Intermediate3 hours