Clean up logic around pmapping the total_energy function.
Clean up logic around pmapping the total_energy function.

Instead of conditionally pmapping the total_energy function in two places
(which, surprisingly, is actually quite easy to trigger both cases, leading to
odd errors), just pmap once non-conditionally.

PiperOrigin-RevId: 407302268
Change-Id: I5ee109be9c3233f110f92203258e879804cb789b

deepmind/ferminet