Fix batching bug

wesselb · Apr 22, 2022 · 18f0e67 · 18f0e67
1 parent cbeabe7
commit 18f0e67
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stheno/random.py b/stheno/random.py
@@ -298,7 +298,7 @@ def kl(self, other: "Normal"):
             scalar: KL divergence.
         """
         return (
-            B.iqf_diag(other.var, other.mean - self.mean)[0]
+            B.iqf_diag(other.var, other.mean - self.mean)[..., 0]
             + B.ratio(self.var, other.var)
             + B.logdet(other.var)
             - B.logdet(self.var)