Dreambooth colab - add option to reset cuda memory

[jmp909]

Is your feature request related to a problem? Please describe.
Sometimes retraining in the same session gives a CUDA memory allocation error

Describe the solution you'd like
suggest potentially adding this as an option to clear memory:

pip install numba

then

from numba import cuda 
device = cuda.get_current_device()
device.reset()

worked for me!

thanks

[ShivamShrirao]

I was thinking of just doing

exit()

Does it work for you or cause any problems ?

[jmp909]

@ShivamShrirao which line are you putting the exit() on? also where would you re-run from?

normally I will just want to run the block to change the MODEL_NAME and OUTPUT_DIR then run the training again to create a an updated copy of the model