Learn.purge() after every TTA iteration

Part 1 (2019) Advanced (Part 1 v3)
1

Hi,
Is it possible to free up some unwanted memory after every iteration of learn.TTA()
I am facing up a problem wherein the GPU memory gets filled up in the 6th iteration and it throws a CUDA memory error. I am using a GTX 1080ti with 11 gigs of memory.

Thanks

2

Maybe create a custom callback and add code to clear up GPU memory inside on_batch_end() link. Try and tell if it works.

3

Thanks for the reply. I am pretty new to fastai, I am not pretty sure whether a callback can be added to the learner at the inference phase. Is there a way we can add a callback at the inference phase as well.

Thanks

4

Callbacks are added when learn is created. Inside that callback add the condition to perform GPU clean up when the model is in eval mode only. learn.model.training can help you with that.