            accelerator.clip_grad_norm_(model.parameters(), max_grad_norm)
            optimizer.step()
            lr_scheduler.step()
            optimizer.zero_grad()
            progress_bar.update(1)

            acc_loss = acc_loss / int(config["gradient_accumulation_steps"])
            accelerator.log({"loss": acc_loss.item()})
            acc_loss = 0
            if completed_steps % checkpointing_steps == 0:
                save_checkpoint(model, accelerator, tokenizer, optimizer, lr_scheduler, save_dir, checkpointing_steps)
