torch.Tensor(x).to(device)
时间: 2024-03-12 18:47:03 浏览: 98
pytorch查看torch.Tensor和model是否在CUDA上的实例
这段代码的作用是将一个形状为 x 的 tensor 转移到指定的 device 上。其中,x 可以是一个列表、元组或者 ndarray,device 可以是 CPU 或者 GPU 设备。这个函数返回一个新的 tensor,并不会改变原有的 tensor。
例如,如果我们想将一个形状为 (3, 4) 的 tensor 转移到 GPU 上,可以使用以下代码:
```
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
x = torch.randn(3, 4)
x_gpu = torch.Tensor(x).to(device)
```
其中,torch.cuda.is_available() 判断当前系统是否支持 GPU,如果支持则使用 GPU 设备,否则使用 CPU 设备。最后得到的 x_gpu 就是一个形状为 (3, 4) 的 tensor,且存储在指定的 device 上。
阅读全文