From 00debdce7615f3e19c2daba1efb5022227686bf0 Mon Sep 17 00:00:00 2001
From: t-werner02 <tobias.werner@fau.de>
Date: Tue, 24 Jun 2025 16:58:28 +0200
Subject: [PATCH] Avoid inappropriate squeeze

---
 src/spainn/plotting.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/spainn/plotting.py b/src/spainn/plotting.py
index 650c758..bc9dcbc 100644
--- a/src/spainn/plotting.py
+++ b/src/spainn/plotting.py
@@ -168,7 +168,10 @@ def _get_data4set(
                     struc[prop if prop != "smooth_nacs" else "nacs"].numpy()
                     for struc in structure
                 ]
-            ).squeeze()
+            )
+
+            if prop == 'energy':
+                ref[prop] = ref[prop].squeeze(axis=1)
 
             if prop in ['smooth_nacs', 'nacs', 'dipoles']:
                 ref[prop], pred[prop] = self._phase_correction(