pytorch多gpu多模型训练时报错AssertionError: Invalid device id

初入pytorch，机器上有两块gpu，分别同时训练两个模型，分别添加了如下代码加以区分：

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

第一块gpu上的模型正常运行，但在第二块gpu上的模型运行时报错 AssertionError: Invalid device id，定位错误代码为：

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

model = DFL_VGG16(k=10, nclass=args.nclass)
device = torch.device("cuda:1")
model = model.to(device)    # 定位到这一行

于是进行了如下测试：

import torch
import os

os.environ["CUDA_VISIBLE_DEVICES"] = "0"    
print(torch.cuda.device_count())            # 1
print(torch.cuda.current_device())          # 0
print(torch.cuda.get_device_name())         # GeForce RTX 2080 Ti

os.environ["CUDA_VISIBLE_DEVICES"] = "1"    
print(torch.cuda.device_count())            # 1
print(torch.cuda.current_device())          # 0
print(torch.cuda.get_device_name())         # GeForce RTX 2080 Ti

即表明如果指定了某一块gpu，不管指定的是第几块，它的名字都为"cuda:0"

所以将错误代码进行如下修改，代码就正常运行了！

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

model = DFL_VGG16(k=10, nclass=args.nclass)
device = torch.device("cuda:0")    # 不管指定的是第几块gpu，gpu的id都为0
model = model.to(device)