当数据和要加载的模型不在同一个cuda上时的处理方法
问题:
c = F.relu(policy_model.Q.c1(exp_list[i][k].unsqueeze(0)))
error:
*** RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cuda:2! (when checking argument for argument weight in
method wrapper__cudnn_convolution)
原因:
训练好的Model和待使用的数据分别在两个cuda上,因此出现错误。
解决方法:
在加载模型时,设置加载到数据所在的cuda上。确保模型和数据在同一个cuda上。
即把:
policy_model = torch.load(PATH)
改为:
policy_model = torch.load(PATH, map_location=torch.device('cuda:1'))