Set device with device argument functions and get it in PyTorch
Tags: pytorch, device, argument, function
Published: 2024-06-25 00:27:46

You can set `device` with the functions which have `device` arguments and get it with [device](https://pytorch.org/docs/stable/tensor_attributes.html#torch-device) as shown below:

*Memos:
- I selected some popular `dtype` argument functions such as [tensor()](https://pytorch.org/docs/stable/generated/torch.tensor.html), [arange()](https://pytorch.org/docs/stable/generated/torch.arange.html), [rand()](https://pytorch.org/docs/stable/generated/torch.rand.html), [rand_like()](https://pytorch.org/docs/stable/generated/torch.rand_like.html), [zeros()](https://pytorch.org/docs/stable/generated/torch.zeros.html) and [zeros_like()](https://pytorch.org/docs/stable/generated/torch.zeros_like.html).
- `device`(Optional-Type:`int`, `str` or [device()](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device)).
- If `device` is not given, the `device` of [set_default_device()](https://pytorch.org/docs/stable/generated/torch.set_default_device.html) 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](https://dev.to/hyperkai/istensor-numel-and-device-in-pytorch-2eha) explains [device()](https://pytorch.org/docs/stable/tensor_attributes.html#torch-device).

`tensor()`. *[My post](https://dev.to/hyperkai/create-a-tensor-in-pytorch-127g) explains `tensor()`:

```python
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()](https://pytorch.org/docs/stable/generated/torch.cuda.is_available.html). *[My post](https://dev.to/hyperkai/check-pytorch-version-cpu-and-gpucuda-in-pytorch-6jk) explains `is_available()`:

```python
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](https://dev.to/hyperkai/arange-linspace-logspace-and-normal-in-pytorch-a87) explains `arange()`:

```python
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](https://dev.to/hyperkai/rand-randlike-randn-randnlike-randint-and-randperm-in-pytorch-31nc) explains `rand()`:

```python
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](https://dev.to/hyperkai/rand-randlike-randn-randnlike-randint-and-randperm-in-pytorch-31nc) explains `rand_like()`:

```python
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](https://dev.to/hyperkai/zeros-zeroslike-ones-and-oneslike-in-pytorch-26jm) explains `zeros()`:

```python
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](https://dev.to/hyperkai/zeros-zeroslike-ones-and-oneslike-in-pytorch-26jm) explains `zeros_like()`:

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