Name: Set device with device argument functions and get it in PyTorch
Rating: 4.1 (9729 reviews)
Author: hyperkai

Buy Me a Coffee☕

*Memos:

My post explains how to set dtype with dtype argument functions and get it.
My post explains how to set requires_grad with requires_grad argument functions and get it.
My post explains how to set keepdim with keepdim argument functions.
My post explains how to set out with out argument functions.

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

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

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