Usage

import torch
x = torch.tensor([1, 2, 3])  # 定義一個變數為 tensor x
x = x.cuda(0)  # 將 x 轉到 GPU(cuda) 上
print(x.device)  # x 在哪個設備上 (cuda or cpu)
x = torch.tensor([1, 2, 3], device='cuda')
x = torch.tensor([1, 2, 3]).to('cuda')
y = x**2
print(y.device)  # 如果對 GPU 變數做運算，結果會在相同的設備上
z = y + x.cpu()  # 不同的設備上的變數不可以直接運算，包括兩個相異 GPU 上也是


import torch
from torch import nn
# 與 tensor 類似，pytorch 模型也可以透過 .cuda() 轉移到設備上。
module = nn.Linear(3, 1)
list(module.parameters())[0].device # device(type='cpu')

module.cuda(0)
list(module.parameters())[0].device  #device=(type='cuda', index=0)

Example

Reference

pytorch使用GPU计算-学习笔记

[开发技巧]·PyTorch如何使用GPU加速（CPU与GPU数据的相互转换）_小宋是呢的博客-CSDN博客_pytorch使用gpu