pytorch: RuntimeError: Cuda error: out of memory

problem¶

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.

reason and solution¶

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:

1. CUDA error: out of memory when load models
2. Out of memory error when resume training even though my GPU is empty
3. SAVING AND LOADING MODELS ACROSS DEVICES IN PYTORCH
4. RunTime Error : cuda out of memory