def _loop_tqdm(func):
        """
        Decorator that adds a tqdm progress bar to `body_fun`
        used in `jax.lax.fori_loop`.
        """
        def wrapper_progress_bar(i, val):
            _update_progress_bar(i)
            result = func(i, val)
            return close_tqdm(result, i)

        return wrapper_progress_bar

        Decorator that adds a tqdm progress bar to `body_fun`
        used in `jax.lax.fori_loop`.
        
cola/utils/jax_tqdm.py