How to move a Torch Tensor from CPU to GPU and vice versa?



A torch tensor defined on CPU can be moved to GPU and vice versa. For high-dimensional tensor computation, the GPU utilizes the power of parallel computing to reduce the compute time.

High-dimensional tensors such as images are highly computation-intensive and takes too much time if run over the CPU. So, we need to move such tensors to GPU.

Syntax

  • To move a torch tensor from CPU to GPU, following syntax/es are used −

Tensor.to("cuda:0") or Tensor.to(cuda)

And,

Tensor.cuda()
  • To move a torch tensor from GPU to CPU, the following syntax/es are used −

Tensor.to("cpu")

And,

Tensor.cpu()

Let's take a couple of examples to demonstrate how a tensor can be moved from CPU to GPU and vice versa.

Note − I have provided two different outputs for each program. One output for the systems having CPU only and the other output for system having GPU along with CPU.

Example 1

# Python program to move a tensor from CPU to GPU # import torch library import torch # create a tensor x = torch.tensor([1.0,2.0,3.0,4.0]) print("Tensor:", x) # check tensor device (cpu/cuda) print("Tensor device:", x.device) # Move tensor from CPU to GPU # check CUDA GPU is available or not print("CUDA GPU:", torch.cuda.is_available()) if torch.cuda.is_available():    x = x.to("cuda:0")    # or x=x.to("cuda") print(x) # now check the tensor device print("Tensor device:", x.device)

Output 1 − When GPU is not available

Tensor: tensor([1., 2., 3., 4.]) Tensor device: cpu CUDA GPU: False tensor([1., 2., 3., 4.]) Tensor device: cpu

Output 2 − When GPU is available

Tensor: tensor([1., 2., 3., 4.]) Tensor device: cpu CUDA GPU: True tensor([1., 2., 3., 4.], device='cuda:0') Tensor device: cuda:0

Example 2

# Python program to move a tensor from CPU to GPU # import torch library import torch # create a tensor on CPU x = torch.tensor([1.0,2.0,3.0,4.0]) print("Tensor:", x) print("Tensor device:", x.device) # Move tensor from CPU to GPU if torch.cuda.is_available():    x = x.cuda() print(x) # now check the tensor device print("Tensor device:", x.device)

Output 1 − If GPU is not available

Tensor: tensor([1., 2., 3., 4.]) Tensor device: cpu tensor([1., 2., 3., 4.]) Tensor device: cpu

Output 2 − If GPU is available

Tensor: tensor([1., 2., 3., 4.]) Tensor device: cpu tensor([1., 2., 3., 4.], device='cuda:0') Tensor device: cuda:0

Example 3

# Python program to move a tensor from GPU to CPU # import torch library import torch # create a tensor on GPU if torch.cuda.is_available():    x = torch.tensor([1.0,2.0,3.0,4.0], device = "cuda") print("Tensor:", x) print("Tensor device:", x.device) # Move tensor from GPU to CPU x = x.to("cpu") # x = x.cpu() print(x) # Now check the tensor device print("Tensor device:", x.device)

Output 1 − If GPU is not available

Tensor: tensor([1., 2., 3., 4.]) Tensor device: cpu tensor([1., 2., 3., 4.]) Tensor device: cpu

Output 2 − If GPU is available

Tensor: tensor([1., 2., 3., 4.], device='cuda:0') Tensor device: cuda:0 tensor([1., 2., 3., 4.]) Tensor device: cpu
Updated on: 2023-11-06T03:42:20+05:30

26K+ Views

Kickstart Your Career

Get certified by completing the course

Get Started
Advertisements