当数据和要加载的模型不在同一个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'))

       

posted @ 2022-02-23 14:52  呦呦南山  阅读(3210)  评论(0编辑  收藏  举报