From 3a0f46dcfc022b4fdc58a4983b658d1bd95200ad Mon Sep 17 00:00:00 2001
From: EnricoTrizio <enricotrizio@gmail.com>
Date: Thu, 9 May 2024 09:53:40 +0200
Subject: [PATCH] Fixed size bug

---
 mlcolvar/utils/plot.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlcolvar/utils/plot.py b/mlcolvar/utils/plot.py
index 1f1326c7..964b422d 100644
--- a/mlcolvar/utils/plot.py
+++ b/mlcolvar/utils/plot.py
@@ -161,7 +161,7 @@ def plot_isolines_2D(
                 with torch.no_grad():
                     train_mode = function.training
                     function.eval()
-                    s = function(xy.unsqueeze(0)).numpy()
+                    s = function(xy.unsqueeze(0)).squeeze(0).numpy()
                     function.training = train_mode
                     if component is not None:
                         s = s[component]