torch.device
时间: 2023-10-17 20:20:49 浏览: 129
torch.device is a class in PyTorch that represents the device on which tensors and neural network models are located. It can be used to specify the device on which operations should be performed, and to move tensors and models between different devices (e.g., between CPU and GPU).
For example, to create a tensor on the CPU, you can use:
```python
import torch
x = torch.tensor([1, 2, 3], device=torch.device('cpu'))
```
To move a tensor to the GPU, you can use:
```python
x = x.to(torch.device('cuda'))
```
Similarly, to create a neural network model on the GPU, you can use:
```python
import torch.nn as nn
model = nn.Linear(10, 1).to(torch.device('cuda'))
```
Note that the device argument can also be set when creating a tensor or a model, e.g. `x = torch.tensor([1, 2, 3], device='cpu')`.