Add config option to specify float precision type during training

facebookresearch / nougat

Implementation of Nougat Neural Optical Understanding for Academic Documents

https://facebookresearch.github.io/nougat/

MIT License

8.98k stars 567 forks source link

Open OrianeN opened 1 year ago

OrianeN commented 1 year ago

Not all CUDA devices support mixed-type precision - myself I got this error when launching train.py:

RuntimeError: Current CUDA Device does not support bfloat16. Please switch dtype to float16.

Therefore I propose this PR to let the user choose which type of precision they want to use during training.