PyTorch默認使用從0開始的GPU，如果GPU0正在運行程序，需要指定其他GPU。

有如下兩種方法來指定需要使用的GPU。

1.類似tensorflow指定GPU的方式，使用CUDA_VISIBLE_DEVICES。

1.1 直接終端中設定：

CUDA_VISIBLE_DEVICES=1 python my_script.py

1.2 python代碼中設定：

import os os.environ["CUDA_VISIBLE_DEVICES"] = "2"

見網址：http://www.cnblogs.com/darkknightzh/p/6591923.html

2. 使用函數 set_device

import torch torch.cuda.set_device(id)

該函數見 pytorch-master\torch\cuda_init_.py。
不過官方建議使用CUDA_VISIBLE_DEVICES，不建議使用 set_device 函數。

該文章為轉載文章，原文章地址:http://www.cnblogs.com/darkknightzh/p/6836568.html

總結

以上是生活随笔為你收集整理的PyTorch中使用指定的GPU的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。