Squeeze predictive mean

ziatdinovmax · Feb 6, 2024 · eca96a2 · eca96a2
1 parent ce93965
commit eca96a2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/gpax/models/sparse_gp.py b/gpax/models/sparse_gp.py
@@ -197,7 +197,7 @@ def get_mvn_posterior(
 
         Linv_W_Dinv_y = Linv_pack[:, :W_Dinv_y.shape[1]]
         Linv_Ws = Linv_pack[:, W_Dinv_y.shape[1]:]
-        mean = Linv_W_Dinv_y.T @ Linv_Ws
+        mean = (Linv_W_Dinv_y.T @ Linv_Ws).squeeze()
 
         Kss = self.kernel(X_new, X_new, params, noise_p)
         Qss = Ws.T @ Ws