Upload model to GCS during multi-TPU training in colab, cell stuck

Hi

I am trying to run multi-TPU training on colab using notebook_launcher. In the training loop, I am able to save the model locally using accelerator.save(model.state_dict(), local_file_name). However when I try to upload it to GCS using gsutil cp, the cell gets stuck

accelerator.print('Inside saving')
accelerator.save(model.state_dict(), local_file_name)
accelerator.print('Model saved')
accelerator.print('Uploading model')
if accelerator.is_main_process:
  gcs_file_name = os.path.join(MODELS_DIR, args.save_prefix, '{}.pt'.format(num_steps))
  ! gsutil cp {local_file_name} {gcs_file_name}
accelerator.print('Upload done to {}!'.format(gcs_file_name))

Output and where it gets stuck:

Inside saving
Model saved
Uploading model

Is it not possible to upload to GCS during train loop? If so, is there another solution?

Thanks

[sgugger]

You could push your model to the Hugging Face Hub when you want to checkpoint (with the push_to_hub method), this way you can resume training from anywhere.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.