Fix energy tests
Fix .cache() to avoid double computation