System Info

Adding a PushToHubCallback callback when training a TF model in a Jupyter notebook results in a cell hanging upon training completion. Nothing is pushed to Hub. Here's the callback:

push_to_hub_callback = PushToHubCallback(
    output_dir="my_food_classifier",
    tokenizer=image_processor,
)
callbacks = [metric_callback, push_to_hub_callback]
model.fit(
    tf_train_dataset,
    validation_data=tf_eval_dataset,
    epochs=num_epochs,
    callbacks=callbacks
)

A Jupyter notebook where this can be reproduced is linked below, however, I'm getting the same result, when running this as a script, not in a notebook environment.

• transformers version: 4.25.1
• Platform: Linux-5.10.147+-x86_64-with-glibc2.27
• Python version: 3.8.16
• Huggingface_hub version: 0.11.1
• PyTorch version (GPU?): 1.13.0+cu116 (True)
• Tensorflow version (GPU?): 2.9.2 (True)
• Flax version (CPU?/GPU?/TPU?): not installed (NA)
• Jax version: not installed
• JaxLib version: not installed
• Using GPU in script?: yes
• Using distributed or parallel set-up in script?: no

Who can help?

@Rocketknight1 @amyeroberts

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

https://colab.research.google.com/drive/1OlFUm41Tqfz4v4oHJ4XzJEpSwmlGq9z3#scrollTo=daKTh8apJHU_

Expected behavior

I would expect that the callback would save and push the model to the Hub once per epoch, and, possibly, upon training completion.