DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on • Edited on

Set and get `device` in PyTorch

Buy Me a Coffee

*Memos:

You can set and get device as shown below:

*Memos:

  • I selected some popular device argument functions such as tensor(), arange(), rand(), rand_like(), zeros() and zeros_like().
  • Basically, device(Optional-Default:None-Type:int, str or device()).
  • Basically, if device is None, it's inferred from other tensor or get_default_device() is used. *My post explains get_default_device() and set_default_device().
  • cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia or privateuseone can be set to device.
  • Setting 0 to device uses cuda(GPU). *The number must be zero or positive.
  • Basically, device= must be needed.
  • My post explains device().
  • str() can get a device value.

tensor(). *My post explains tensor():

import torch

my_tensor = torch.tensor([0, 1, 2])
my_tensor = torch.tensor([0, 1, 2], device='cpu')
my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cpu'))
my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cpu'))

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0, 1, 2]), device(type='cpu'), 'cpu')

my_tensor = torch.tensor([0, 1, 2], device='cuda:0')
my_tensor = torch.tensor([0, 1, 2], device='cuda')
my_tensor = torch.tensor([0, 1, 2], device=0)
my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cuda:0'))
my_tensor = torch.tensor([0, 1, 2], device=torch.device(device='cuda'))
my_tensor = torch.tensor([0, 1, 2], device=torch.device(device=0))
my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cuda', index=0))
my_tensor = torch.tensor([0, 1, 2], device=torch.device(type='cuda'))

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=0), 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

tensor() with is_available(). *My post explains is_available():

import torch

my_device = "cuda:0" if torch.cuda.is_available() else "cpu"
my_tensor = torch.tensor([0, 1, 2], device=my_device)

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=0), 'cuda:0')
Enter fullscreen mode Exit fullscreen mode

arange(). *My post explains arange():

import torch

my_tensor = torch.arange(start=5, end=15, step=3, device='cpu')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([5, 8, 11, 14]), device(type='cpu'), 'cpu')

my_tensor = torch.arange(start=5, end=15, step=3, device='cuda:0')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([5, 8, 11, 14], device='cuda:0'),
#  device(type='cuda', index=0),
#  'cuda:0')
Enter fullscreen mode Exit fullscreen mode

rand(). *My post explains rand():

import torch

my_tensor = torch.rand(size=(3,), device='cpu')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0.2782, 0.3780, 0.6509]), device(type='cpu'), 'cpu')

my_tensor = torch.rand(size=(3,), device='cuda:0')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0.1052, 0.9281, 0.0151], device='cuda:0'),
#  device(type='cuda', index=0),
#  'cuda:0')
Enter fullscreen mode Exit fullscreen mode

rand_like(). *My post explains rand_like():

import torch

my_tensor = torch.rand_like(input=torch.tensor([7., 4., 5.]), 
                            device='cpu')
my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0.9130, 0.7072, 0.1935]), device(type='cpu'), 'cpu')

my_tensor = torch.rand_like(input=torch.tensor([7., 4., 5.]), 
                            device='cuda:0')
my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0.3655, 0.6319, 0.3045], device='cuda:0'),
#  device(type='cuda', index=0),
#  'cuda:0')
Enter fullscreen mode Exit fullscreen mode

zeros(). *My post explains zeros():

import torch

my_tensor = torch.zeros(size=(3,), device='cpu')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0., 0., 0.]), device(type='cpu'), 'cpu')

my_tensor = torch.zeros(size=(3,), device='cuda:0')

my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0., 0., 0.], device='cuda:0'),
#  device(type='cuda', index=0),
#  'cuda:0')
Enter fullscreen mode Exit fullscreen mode

zeros_like(). *My post explains zeros_like():

import torch

my_tensor = torch.zeros_like(input=torch.tensor([7., 4., 5.]), 
                             device='cpu')
my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0., 0., 0.]), device(type='cpu'), 'cpu')

my_tensor = torch.zeros_like(input=torch.tensor([7., 4., 5.]), 
                             device='cuda:0')
my_tensor, my_tensor.device, str(my_tensor.device)
# (tensor([0., 0., 0.], device='cuda:0'),
#  device(type='cuda', index=0),
#  'cuda:0')
Enter fullscreen mode Exit fullscreen mode

Image of Docusign

🛠️ Bring your solution into Docusign. Reach over 1.6M customers.

Docusign is now extensible. Overcome challenges with disconnected products and inaccessible data by bringing your solutions into Docusign and publishing to 1.6M customers in the App Center.

Learn more

Top comments (0)

A Workflow Copilot. Tailored to You.

Pieces.app image

Our desktop app, with its intelligent copilot, streamlines coding by generating snippets, extracting code from screenshots, and accelerating problem-solving.

Read the docs