成功解决:RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument for argument index in method wrapper__index_select),报错问题。
问题原因:模型训练时,数据不在同一块GPU上
解决方案:在模型训练的代码最前面加上下面的代码即可
import os
# 指定使用一块GPU,[0,1,2,3]下面是只选择使用1号GPU
os.environ["CUDA_VISIBLE_DEVICES"] = "1"