implement training on per atom energies by modifying loss modules

implement training on per atom energies by modifying loss modules #523