From 42cae0aaddf737a3b2506046a89ec86c18850667 Mon Sep 17 00:00:00 2001
From: Daniel Dodd <daniel_dodd@icloud.com>
Date: Sun, 8 Jan 2023 21:18:18 +0000
Subject: [PATCH] Update variational_inference.py

---
 gpjax/variational_inference.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gpjax/variational_inference.py b/gpjax/variational_inference.py
index 6745656e..7dc749bd 100644
--- a/gpjax/variational_inference.py
+++ b/gpjax/variational_inference.py
@@ -142,7 +142,7 @@ def q_moments(x):
         log_prob = vmap(lambda f, y: link_function(params["likelihood"], f).log_prob(y))
 
         # ≈ ∫[log(p(y|f(x))) q(f(x))] df(x)
-        expectation = gauss_hermite_quadrature(log_prob, mean, variance, y=y)
+        expectation = gauss_hermite_quadrature(log_prob, mean, jnp.sqrt(variance), y=y)
 
         return expectation