PushToHubCallback is hanging on the training completion
Author: MKhalusova
Pull request: no

### 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
- [X] My own modified scripts

### Tasks

- [X] 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.  

Comments:
`PushToHubCallback(..., save_strategy="no")` seems to fix the hanging problem but in that case the model will only be saved at the end of the training...

Investigating this now - PushToHubCallback shouldn't be hanging like this.

I tried running this locally and it ran successfully for me - is it possible that model uploading from Colab was just slow, and you didn't give it time to complete? In some cases the progress bar on the final upload doesn't display, though I have some ideas on how to fix that!

Hi @Rocketknight1 , thanks for looking into the issue! When I reported this, I gave it about 40-45 minutes after it finished training. At that point, nothing was uploaded to Hub and the notebook cell just kept spinning with no end and I interrupted the execution. However, I have tried to reproduce it again just now, and… it worked. Checkpoints uploaded during training, and final model took less than 2 minutes to upload. I have not modified anything in the code. No clue why it wasn’t working before, and why it does now.  
My only guess is that perhaps there was some sort of connection issue. Could it be that the callback was not able to reach Hub at that moment and kept trying?  

As a piece of anecdotal evidence - in case it's useful - I tried running @MKhalusova example a few days ago (added full script below), and also had issues with hanging. It ran seemlessly today - no idea what changed 🤷‍♀️ 

Two notes: 
* [The model card only](https://huggingface.co/amyeroberts/my_food_classifier) shows performance for epoch 0, even though the save strategy is "epoch". I'm guessing this is because of [not starting an upload job whilst another is happening](https://github.com/huggingface/transformers/blob/91c2278b97a16e7dcde28fd0fce72969560f587b/src/transformers/keras_callbacks.py#L374)? 
* I've never seen a progress bar for upload either at the end of an epoch or at the end of training using `PushToHubCallback` 

```python
import numpy as np
import tensorflow as tf

import evaluate
from datasets import load_dataset
from transformers import AutoImageProcessor, DefaultDataCollator, TFAutoModelForImageClassification, create_optimizer
from transformers.keras_callbacks import KerasMetricCallback, PushToHubCallback

BATCH_SIZE = 16
NUM_BATCHES = 5
NUM_EPOCHS = 5
LEARNING_RATE = 3e-5
WEIGHT_DECAY_RATE = 0.01
SEED = 42

# Load in the dataset
food = load_dataset("food101", split="train[:5000]")
food = food.train_test_split(test_size=0.2)

labels = food["train"].features["label"].names
label2id, id2label = dict(), dict()
for i, label in enumerate(labels):
    label2id[label] = str(i)
    id2label[str(i)] = label

checkpoint = "google/vit-base-patch16-224-in21k"
image_processor = AutoImageProcessor.from_pretrained(checkpoint)
accuracy = evaluate.load("accuracy")

def process(examples):
    examples.update(image_processor(examples['image'], ))
    return examples

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return accuracy.compute(predictions=predictions, references=labels)

food = food.map(process, batched=True).shuffle(seed=SEED)
data_collator = DefaultDataCollator(return_tensors="tf")
model = TFAutoModelForImageClassification.from_pretrained(
    checkpoint,
    id2label=id2label,
    label2id=label2id,
)
tf_train_dataset = food["train"].select(range(BATCH_SIZE * NUM_BATCHES)).to_tf_dataset(
   columns=['pixel_values'],
   label_cols=["label"],
   shuffle=True,
   batch_size=BATCH_SIZE,
   collate_fn=data_collator
)
tf_eval_dataset = food["test"].select(range(BATCH_SIZE * NUM_BATCHES)).to_tf_dataset(
   columns=['pixel_values'],
   label_cols=["label"],
   shuffle=True,
   batch_size=BATCH_SIZE,
   collate_fn=data_collator
)
optimizer, lr_schedule = create_optimizer(
    init_lr=LEARNING_RATE,
    num_train_steps=len(food["train"]) * NUM_EPOCHS,
    weight_decay_rate=WEIGHT_DECAY_RATE,
    num_warmup_steps=0,
)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer=optimizer, loss=loss)
metric_callback = KerasMetricCallback(metric_fn=compute_metrics, eval_dataset=tf_eval_dataset)
push_to_hub_callback = PushToHubCallback(
    output_dir="amyeroberts/my_food_classifier",
    tokenizer=image_processor,
)
model.fit(
    tf_train_dataset,
    validation_data=tf_eval_dataset,
    epochs=NUM_EPOCHS,
    callbacks=[metric_callback, push_to_hub_callback]
)
```

@amyeroberts We generally avoid displaying progress bars from that callback, because the upload of each checkpoint runs in the background while the next epoch is training. As a result, the callback progress bar would run into the Keras progress bar and cause chaos in the console output.

However, I think the callback was supposed to display a progress bar for the final upload after training is finished, when there's no risk of running into the Keras bars. This is also the only upload that will actually cause any delays. I'll put that on my list to investigate!

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.