1. 程式人生 > >PyTorch學習系列(十六)——如何使用cuda進行訓練?

PyTorch學習系列(十六)——如何使用cuda進行訓練?

如果想在CUDA上進行計算,需要將操作物件放在GPU記憶體中。
對於普通的張量,可以直接:

x = torch.randn(2, 3)
x = x.cuda()

對於神經網路:

model = MyModel()
model.cuda()

同一個GPU上的張量計算結果仍然儲存在該GPU上。

參考