Pytorch to（device）

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)

這兩行代碼放在讀取數據之前。

mytensor = my_tensor.to(device)

這行代碼的意思是將所有最開始讀取數據時的tensor變量copy一份到device所指定的GPU上去，之後的運算都在GPU上進行。

這句話需要寫的次數等於需要保存GPU上的tensor變量的個數；一般情況下這些tensor變量都是最開始讀數據時的tensor變量，後面衍生的變量自然也都在GPU上

如果是多個GPU

在代碼中的使用方法爲：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = Model()

if torch.cuda.device_count() > 1:

  model = nn.DataParallel(model，device_ids=[0,1,2])

 

model.to(device)