diff --git a/SU2_CFD/include/variables/CPrimitiveIndices.hpp b/SU2_CFD/include/variables/CPrimitiveIndices.hpp
index b5bfee928f18..97687547d394 100644
--- a/SU2_CFD/include/variables/CPrimitiveIndices.hpp
+++ b/SU2_CFD/include/variables/CPrimitiveIndices.hpp
@@ -124,6 +124,9 @@ std::map<std::string, IndexType> PrimitiveNameToIndexMap(const CPrimitiveIndices
   nameToIndex["VELOCITY_Y"] = idx.Velocity() + 1;
   if (idx.NDim() == 3) {
     nameToIndex["VELOCITY_Z"] = idx.Velocity() + 2;
+    nameToIndex["ENERGY"] = idx.Velocity() + 3;
+  } else {
+    nameToIndex["ENERGY"] = idx.Velocity() + 2;
   }
   for (IndexType iSpecies = 0; iSpecies < idx.NSpecies(); ++iSpecies) {
     nameToIndex["DENSITY_" + std::to_string(iSpecies)] = idx.SpeciesDensities() + iSpecies;
diff --git a/TestCases/py_wrapper/turbulent_premixed_psi/run.py b/TestCases/py_wrapper/turbulent_premixed_psi/run.py
index dc1133169a24..9ed534240ace 100644
--- a/TestCases/py_wrapper/turbulent_premixed_psi/run.py
+++ b/TestCases/py_wrapper/turbulent_premixed_psi/run.py
@@ -100,9 +100,9 @@ def update_temperature(SU2Driver, iPoint):
     iCp = primindex.get("CP_TOTAL")
     Cp = SU2Driver.Primitives()(iPoint,iCp)
     # get solution index for energy equation
-    iTEMP = solindex.get("TEMPERATURE")
-    # set enthalpy, for Ideal Gas model, enthalpy is Cp*T
-    SU2Driver.Solution(iFLOWSOLVER).Set(iPoint,iTEMP,Cp*T)
+    iEnergy = solindex.get("ENERGY")
+    # set solution in energy equation, for Ideal Gas model, enthalpy is Cp*T
+    SU2Driver.Solution(iFLOWSOLVER).Set(iPoint,iEnergy,Cp*T)
 
 
 # ################################################################## #