Fix example in accumulate method (#1211)

huggingface · Mar 19, 2023 · 82c2665 · 82c2665
1 parent 2930cac
commit 82c2665
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -841,8 +841,8 @@ def accumulate(self, model):
         >>> accelerator = Accelerator(gradient_accumulation_steps=1)
         >>> dataloader, model, optimizer, scheduler = accelerator.prepare(dataloader, model, optimizer, scheduler)
 
-        >>> with accelerator.accumulate(model):
-        ...     for input, output in dataloader:
+        >>> for input, output in dataloader:
+        ...     with accelerator.accumulate(model):
         ...         outputs = model(input)
         ...         loss = loss_func(outputs)
         ...         loss.backward()