Mixed precision output

[ilyes319]

Is your feature request related to a problem? Please describe.
The sum over atoms after the readout to predict the total energy is now performed in float32 when the model precision is float32. This leads to a potential large accumulation of errors.

Describe the solution you'd like
Upcast the output to float64 before the sum, and use the E0 as float64 in order to save the most precision.

Describe alternatives you've considered
Using float64 for all the model but more expensive.