Set and get `device` in PyTorch

#python #pytorch #device #argument

Buy Me a Coffee☕

*Memos:

My post explains how to set and get dtype.
My post explains how to set requires_grad and get grad.
My post explains keepdim argument.
My post explains out argument.
My post explains bias argument.

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

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

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

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

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

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

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

DEV Community

Set and get `device` in PyTorch

Top comments (0)

Read next

🚀 Building a User Management API with FastAPI and SQLite

ChatsAPI — The World’s Fastest AI Agent Framework

EchoAPI vs Bruno: A Comprehensive Comparison from Design to Testing 💡

How to Create Your Own RAG with Free LLM Models and a Knowledge Base