关于torch.load()更改了cuda位置还是cuda内存不够的问题

问题背景：在一次任务中，由于需要使用cuda进行代码运行，但是分明修改了cuda到一个空闲的卡位置，但是依然抱错cuda out of memory的问题

在任务中，最开始原始代码是采用以下方式尽心模型load的（指定了cuda to device的位置是空闲的卡2）：
在这里插入图片描述
但是依然抱错cuda out of memory，分析如下：

model.load_state_dict(torch.load(FLAGS.restore))
print(f"Model loaded from {FLAGS.restore}")
model.to(FLAGS.device)

从上面的代码可以看出，其实model是先load到默认路径，然后再转换到了指定的device的位置，这样就在第一次load的时候，如果默认位置的cuda被占用，就会导致out of memory的问题

建议好的习惯：
model.load_state_dict(torch.load(FLAGS.restore, map_location=‘cuda:2’))
在model load 的时候直接map_location到指定位置，例如你可以在config文件中设置cuda卡配置：

model.load_state_dict(torch.load(FLAGS.restore, map_location=FLAGS.device))

这样则免除了load到默认卡的这一步。

respect！！！