Reproducible colab notebook : https://colab.research.google.com/drive/1kLpn88HSn7WOnH5vN5l-GlLRTPVArsDw?usp=sharing
I want to use mixed precision to speed up the trai