Set `device` with `device` argument functions and get it in PyTorch – ██FR█████ █INTELL███████████

This content originally appeared on DEV Community and was authored by Super Kai (Kazuya Ito)

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(int, str or torch.device) (Optional). *Memos:
- 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 uses GPU(CUDA).
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(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))

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))

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))

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))

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))

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))

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))

This content originally appeared on DEV Community and was authored by Super Kai (Kazuya Ito)