CUDA out of memory error

[natedailey7]

After training my model for many epochs (200), I receive the following error when I run learner.save_model_bundle():

OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 4.00 GiB total capacity; 3.43 GiB already allocated; 0 bytes free; 3.49 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

I've done some research on how to clear Cuda's memory (Stack Overflow posts here and here). So far, I've tried running:

torch.no_grad():
    torch.cuda.empty_cache()

And

import gc
gc.collect()

I suspect that my learner variable is the culprit for taking up memory, however Jupyter is not displaying a size, so I'm not sure.

It seems that performing garbage collection and emptying the cache will not work because my learner is taking up so much memory, and I would need to delete it . However, I do not want to delete it (or restart my kernel) because I haven't saved it. Wondering if anyone has encountered something similar and has a solution. Perhaps there is a way to specify no_grad() within Raster Vision?

[AdeelH]

I suspect the ONNX export is the problem (specifically this line that tries to allocate a sample input on the device). You can try learner.save_model_bundle(export_onnx=False) instead.

[natedailey7]

That worked perfectly! Thank you so much!