When loading the trained model for testing, I encountered RuntimeError: Cuda error: out of memory.
I was surprised, because the model is not too big, so the video memory is exploding.
Later, I found the answer on the pytorch forum. It turned out that when loading the model, you need to load it to the cpu
through the map_location
parameter of torch.load()
, and then put it on the gpu
.
def load_model(model, model_save_path, use_state_dict=True):
print(f"[i] load model from {model_save_path}")
device = torch.device("cpu") # Load to cpu first
if use_state_dict:
model.load_state_dict(torch.load(model_save_path, map_location=device))
else:
model = torch.load(model_save_path, map_location=device)
print("[i] done")
return model
model = Model(config)
model = load_model(model, config["model_save_path"], use_state_dict=True)
model.to(config["device"]) # put on gpu
But the official doesn't seem to say the reason for doing this. . For details, please refer to: