pytorch to_device
时间: 2024-08-12 09:04:40 浏览: 42
pytorch神经网络.zip
`to_device`是PyTorch库中的一个功能,主要用于将Tensor(张量)从当前设备移动到指定的设备。它属于`torch.Tensor`类的一个方法,用于加快数据传输和计算效率。例如,如果你有一个Tensor在CPU上,而你想要将其移动到GPU上运行,就可以使用`to_device`:
```python
import torch
# 假设x是一个在CPU上的Tensor
x = torch.tensor([1, 2, 3]).cpu()
# 将x移动到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 检查是否有可用的GPU
x_gpu = x.to(device)
print(x_gpu.device) # 输出 "cuda" 或者 "cpu"
```
在这个例子中,`to_device`可以根据当前是否能使用GPU自动选择设备,如果可以,它会将Tensor移动到GPU,否则就保持在CPU上。
阅读全文