How to checkpoint TFAutoModelForSequenceClassification every k batches
I am following the Huggingface Tensorflow Notebook to [train a BERT model on MNLI](https://github.com/huggingface/notebooks/blob/main/examples/text_classification-tf.ipynb) (the task itself is irrelevant). I am interested in checkpointing models locally every k batches, so that I end with 20 or 30 intermediate checkpoints over the course of training. The notebook uses model.fit() to train the model, but my understanding is that using a [callback](https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint) to save intermediate checkpoints saves a SavedModel file. 

I am currently not saving intermediate checkpoints, only the final model, and use model.save_pretrained(ckpt). However, I now want the model to be checkpointed as a part of training. If I try loading a checkpoint directory in the SavedModel format using `model = tf.keras.models.load_model()`, the expected input format is not compatible with how I feed in inputs for `model = TFAutoModelForSequenceClassification.from_pretrained()`. I get the following message (32 is batch size, I think 109 is the input token length for the longest example in the batch). 
<img width="686" alt="Screen Shot 2022-06-17 at 7 20 14 PM" src="https://user-images.githubusercontent.com/60128552/174418985-4331fca5-4b1b-4486-90f2-4f98b98a35d1.png">


My question is how to checkpoint a `TFAutoModelForSequenceClassification` model with a specified frequency (every k batches). Either I have to change how I am feeding the input data into `model = tf.keras.models.load_model()` or change how I am saving the intermediate checkpoints, but I am having trouble figuring out what to do and how to proceed. 
Note: I am not interested in pushing models to the hub, only saving models locally. 
Hey @preethiseshadri518 👋  We have an outstanding set of issues related to the `SavedModel` format, resulting in the errors you see. There are ways to work around it, by manually specifying portions of the graph at save time -- you can find them if you search in closed issues here :)

A much simpler workaround, if it suits your problem, is to store/load the weights with [`save_weights`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#save_weights) and [`load_weights`](https://www.tensorflow.org/api_docs/python/tf/keras/Model#load_weights), which is what our API uses. Let us know if it sorted your problem!<|||||>Hi Joao! Yes, this solves the issue. If I use `save_weights` and save them in the h5 format, then I can use  `TFAutoModelForSequenceClassification.from_pretrained()`. Thanks for your comment!