【记录】Linux系统 pytorch使用其他显卡的方法

解决即使指定显卡，Pytorch仍在cuda:0上训练出现内存不够的问题。

本篇主要详述用这个方法

import os

os.environ['CUDA_VISIBLE_DEVICES'] = '1,2' #'1'

引号里的数字代表显卡号。

这个方法有2个很重要的点：

1. 需要把代码写在最前面，尤其是要在import torch前面
2. 无论之前显卡的id是什么，在指定了显卡的数量后，都从0开始计数。要及时更新后续的代码段。

                例如：指定id为‘2，4’显卡，那么后文代码中第一张显卡为‘cuda:0’，第二张显卡为‘cuda:1’；

                再如：指定id为‘3’的显卡，后文代码中的显卡号为‘cuda:0’.