Update runners.py

brainpy · Aug 3, 2024 · 21068ff · 21068ff
1 parent 5e75f78
commit 21068ff
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/brainpy/_src/runners.py b/brainpy/_src/runners.py
@@ -631,7 +631,7 @@ def _step_func_predict(self, i, *x, shared_args=None):
 
     # finally
     if self.progress_bar:
-      jax.pure_callback(lambda: self._pbar.update(), ())
+      jax.debug.callback(lambda *args: self._pbar.update(), ())
     # share.clear_shargs()
     clear_input(self.target)