Add support for models with lower precision (like 16-bit floats and 8-bit integers)

[iver56]

Check out https://pypi.org/project/diffq/

[iver56]

Also try precision=16 in pytorch-lightning. By default it uses the native amp in pytorch