saving intermediate checkpoints

[1190303125]

Hello, how can we save intermediate checkpoints with your code?

We add the following code. However, it does not work.

if completed_steps+1%32==0:
  if  accelerator.is_main_process:
      accelerator.print(f"Saving model to {args.output_dir}")
      unwrapped_model = accelerator.unwrap_model(model)
      state_dict = accelerator.get_state_dict(model)
      print(completed_steps+1)
      unwrapped_model.save_pretrained(
          f"{args.output_dir}_{completed_steps}",
          is_main_process=accelerator.is_main_process,
          save_function=accelerator.save,
          state_dict=state_dict,  )
  accelerator.wait_for_everyone()

We find that the program keeps getting stuck at the step state_dict = accelerator.get_state_dict(unwrapped_model).