pytorch指定GPU来训练模型

当前做的实验需要单独使用四块卡中的一个，我原本的代码是四块卡并行计算，所以需要指定GPU来训练模型：
（1）在代码里指定单块GPU的命令如下：
import os
torch.cuda.set_device(3)

（2）在终端命令行指定某一块或某几块GPU：

CUDA_VISIBLE_DEVICES = 0, 1, 2, 3 python train.py或CUDA _VISIBLE_DEVICES=3 python train.py。