accelerator_accumulate

accelerator_accumulate.py

total_loss = 0.0
for batch in dataloader:
    if config.is_comparison:
        batch = get_model_specific_batch(batch, model_name)

    with accelerator.accumulate(model):
        outputs = model(**batch)
        loss = outputs.loss
        total_loss += loss.detach().float()
        accelerator.backward(loss)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    global_step += 1
return total_loss