From 3b231d5daeea80f879b57f3a5f0499003cffd6b7 Mon Sep 17 00:00:00 2001
From: ClaudioZeni <claudiopuntozeni@gmail.com>
Date: Mon, 10 May 2021 16:50:43 +0200
Subject: [PATCH] Bugfix

---
 raffy/linear_potential.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/raffy/linear_potential.py b/raffy/linear_potential.py
index 06899f7..0718c7c 100644
--- a/raffy/linear_potential.py
+++ b/raffy/linear_potential.py
@@ -65,7 +65,7 @@ def add_square_g(self, g, dg, X, compute_forces):
 
     def adjust_g(self, g, dg, X, compute_forces=True, train_pca=False):
         dg_reshape = []
-        if len(g.shape) == 3:
+        if len(g[0].shape) == 2:
             g = np.array([x.sum(axis=0) for x in g])
         for i in np.arange(len(g)):
             if compute_forces and dg[i].shape == 4: