DEV Community

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

Posted on • Updated on

Set device with device argument functions and get it in PyTorch

#pytorch #device #argument #function

You can set device with the functions which have device arguments and get it with device as shown below:

*Memos:

  • I selected some popular dtype argument functions such as tensor(), arange(), rand(), rand_like(), zeros() and zeros_like().
  • device(Optional-Type:int, str or device()).
  • If device is not given, the device of set_default_device() is used.
  • 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= is needed.
  • My post explains device().

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
# (tensor([0, 1, 2]), device(type='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
# (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=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
# (tensor([0, 1, 2], device='cuda:0'), device(type='cuda', index=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
# (tensor([5, 8, 11, 14]), device(type='cpu'))

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

my_tensor, my_tensor.device
# (tensor([5, 8, 11, 14], device='cuda:0'), device(type='cuda', index=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
# (tensor([0.2985, 0.4517, 0.1018]), device(type='cpu'))

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

my_tensor, my_tensor.device
# (tensor([0.6161, 0.8663, 0.8344], device='cuda:0'),
#  device(type='cuda', index=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
# (tensor([0.8479, 0.3738, 0.7446]), device(type='cpu'))

my_tensor = torch.rand_like(input=torch.tensor([7., 4., 5.]), 
                            device='cuda:0')
my_tensor, my_tensor.device
# (tensor([0.2788, 0.1682, 0.3529], device='cuda:0'),
#  device(type='cuda', index=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
# (tensor([0., 0., 0.]), device(type='cpu'))

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

my_tensor, my_tensor.device
# (tensor([0., 0., 0.], device='cuda:0'), device(type='cuda', index=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
# (tensor([0., 0., 0.]), device(type='cpu'))

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

Top comments (0)

Read next

blog_ts profile image

The Synergy Between Design and Marketing: Crafting a Cohesive Brand Experience

Sara Tanning -

webfaisalbd profile image

in product details page, go to the top position on the details page

Faisal Ahmed -

elitenoire profile image

Styling React.js UIs

Eva Raymond -

land-bit profile image

Overcoming the challenges of backend development : My experience and application for the HNG internship

Landry Bitege -