pytorch 指定GPU训练

# 1：
torch.cuda.set_device(1)

# 2：
device = torch.device("cuda:1")

# 3：（官方推荐）
import os
os.environ["CUDA_VISIBLE_DEVICES"] = '1'

（同时调用两块GPU的话）
os.environ["CUDA_VISIBLE_DEVICES"] = '1,2'