Half Precision implemented for Low VRAM GPUs

[cperales]

Hi there!!

First of all, amazing job with your library. It achieves great results, and it is not difficult to implement.

I have several laptops, and I could run this and DeepDaze on my new laptop, but I couldn't run BigSleep on my old computer, with a GPU GTX 1050 with 4 GB of VRAM.

So I decided to implement a boolean parameter that reduces the precision of the model in the train_step method. Besides, I added another parameter, image_folder. This parameter can be a string, naming the folder where to save the images.

[lucidrains]

@cperales oh hey Carlos! this looks great! ❤️ do you want to try extending this to the CLI as well?

[cperales]

Well, I am not so sure about how to do it... Is it enough if I modifed the file cli.py, adding the options to train function?

[cperales]

Btw, I found out that, better that some code can be simplified with context with torch.cuda.amp.autocast():.

https://spell.ml/blog/mixed-precision-training-with-pytorch-Xuk7YBEAACAASJam