        "train_loss": train_loss.item(),
        "epoch": epoch,
        "step": completed_steps,
    }
    accelerator.log(base_log, step=completed_steps)

    import datetime
    timeout = datetime.timedelta(seconds=100000000)
    accelerator = Accelerator(
        gradient_accumulation_steps=int(config["gradient_accumulation_steps"]),
        mixed_precision="no",
