Resolve issue that broke tests for returning var in spacetime

AaltoML · Nov 21, 2022 · 307d9e0 · 307d9e0
1 parent af3f625
commit 307d9e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bayesnewton/basemodels.py b/bayesnewton/basemodels.py
@@ -813,7 +813,7 @@ def predict(self, X=None, R=None, pseudo_lik_params=None):
 
         # if np.squeeze(test_var).ndim > 2:  # deal with spatio-temporal case (discard spatial covariance)
         if self.spatio_temporal:  # deal with spatio-temporal case (discard spatial covariance)
-            test_var = diag(np.squeeze(test_var))
+            test_var = diag(test_var)
         return np.squeeze(test_mean), np.squeeze(test_var)
 
     def filter_energy(self):