lint the code

awslabs · Dec 19, 2023 · 4c4bd0b · 4c4bd0b
1 parent 4d9dca9
commit 4c4bd0b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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: