CUDA out of memory when calling model.state_dict()

[bily]

I have been training a very large model which occupies over 90% GPU memory. The training process works fine until calling model.state_dict() to save checkpoints. It throws CUDA out of memory error.

The same problem has been reported in:

Why CUDA runs out of memory when calling torch.save?

+1. I had the same problem and moving the model to cpu before saving also worked.

Out of memory during torch.save

Shifting the model to CPU before saving and moving it back to GPU after saving solved it.

The suggested fix is to move the model to CPU device before saving checkpoints and then move the model back to GPU. However, I use distributed training and the training program freezes after model.cuda().

Any help is appreciated!