Out of memory during torch.save

2

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

4 Likes