src/transformers/integrations/integration_utils.py/FlyteCallback/on_save
class FlyteCallback:    def on_save(self, args, state, control, **kwargs):
        if self.sync_checkpoints and state.is_world_process_zero:
            ckpt_dir = f"checkpoint-{state.global_step}"
            artifact_path = os.path.join(args.output_dir, ckpt_dir)

            logger.info(f"Syncing checkpoint in {ckpt_dir} to Flyte. This may take time.")
            self.cp.save(artifact_path)