diff --git a/pes_to_spec/model.py b/pes_to_spec/model.py
index ce57af8a1520c1f27071dda5ce4c41cf3259ab0c..e8fe9eff8ed722e59b1e6b3c17567bab1ee82967 100644
--- a/pes_to_spec/model.py
+++ b/pes_to_spec/model.py
@@ -504,10 +504,10 @@ class SelectRelevantLowResolution(TransformerMixin, BaseEstimator):
                  for channel, item in X.items()}
         if not keep_dictionary_structure:
             selected = list(y.values())
-            if self.poly:
-                selected += [np.sqrt(np.fabs(v)) for v in y.values()]
             if pulse_energy is not None:
                 selected += [pulse_energy[:, np.newaxis, :]]
+                if self.poly:
+                    selected += [pulse_energy[:, np.newaxis, :]*v for v in y.values()]
             return np.concatenate(selected, axis=-1)
         return y