##
from accelerate import Accelerator
accelerator = Accelerator()
dataloader, model, optimizer, scheduler = accelerator.prepare(
dataloader, model, optimizer, scheduler
)
for batch in dataloader:
inputs, targets = batch
outputs = model(inputs)
loss = loss_function(outputs, targets)
accelerator.backward(loss)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
+accelerator.save_state("checkpoint_dir")
+accelerator.load_state("checkpoint_dir")
##
To save or load a checkpoint in, `Accelerator` provides the `save_state` and `load_state` methods.
These methods will save or load the state of the model, optimizer, scheduler, as well as random states and
any custom registered objects from the main process on each device to a passed in folder.
**This API is designed to save and resume training states only from within the same python script or training setup.**
##
To learn more checkout the related documentation:
- Saving and loading training states
- `save_state` API reference
- `load_state` API reference
- Example script