From 18f0e678520f84dba0443b4ba4e5c60c08353084 Mon Sep 17 00:00:00 2001
From: Wessel Bruinsma <wessel.p.bruinsma@gmail.com>
Date: Fri, 22 Apr 2022 09:07:10 +0200
Subject: [PATCH] Fix batching bug

---
 stheno/random.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stheno/random.py b/stheno/random.py
index 71a0226..9f171be 100644
--- 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)