From 4c4bd0bad31766478c00c010c3759c9a066c45d8 Mon Sep 17 00:00:00 2001
From: Gianluca Detommaso <detommaso.gianluca@gmail.com>
Date: Tue, 19 Dec 2023 12:07:42 +0100
Subject: [PATCH] lint the code

---
 fortuna/prob_model/predictive/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/fortuna/prob_model/predictive/base.py b/fortuna/prob_model/predictive/base.py
index 3d87387c..4218a775 100644
--- a/fortuna/prob_model/predictive/base.py
+++ b/fortuna/prob_model/predictive/base.py
@@ -564,7 +564,7 @@ def _sample(key, _inputs):
             if distribute:
                 outputs = jnp.stack(
                     list(map(lambda key: _sample(shard_prng_key(key), inputs), keys)),
-                    axis=1
+                    axis=1,
                 )
                 outputs = self._unshard_ensemble_arrays(outputs)
             else: