AssertionError: Invalid device id

原代码：

device = th.device("cuda:2" if th.cuda.is_available() else "cpu")
device_ids = [2, 1, 3]

报错：

AssertionError: Invalid device id

原因：

造成这个错误的原因是GPU编号原因
pytorch默认使用gpu编号为device:0的设备
本地只有一个 GPU (GPU:0)，所以编号只有0

修改为：

device = th.device("cuda:0" if th.cuda.is_available() else "cpu")
device_ids = [0]

有的没的链接：

pytorch关于多块gpu使用总结，报错AssertionError: Invalid device id

Runtimeerror: cuda error: invalid device ordinal

pytorch分布式训练（一）：torch.nn.DataParallel

AssertionError: Invalid device id