def save_ckpt(ckpt_dir: str, commit_msg: str = ""):
        """save checkpoints and push to Hugging Face Hub if specified"""

        # save checkpoint after each epoch and push checkpoint to the hub
        if jax.process_index() == 0:
            params = jax.device_get(jax.tree_util.tree_map(lambda x: x[0], state.params))
            model.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir), params=params)
            tokenizer.save_pretrained(os.path.join(training_args.output_dir, ckpt_dir))
            if training_args.push_to_hub:
                repo.push_to_hub(commit_message=commit_msg, blocking=False)
save checkpoints and push to Hugging Face Hub if specified
dpkd/transformers/examples/flax/image-captioning/run_image_captioning_flax.py