Segformer, can't save checkpoint in saved_model format
Thanks for this repo ! 
###  Issue description:
1) initalise your segformer model
2) try to save it.

### Error:
```
    File "/home/maxime/anaconda3/envs/tf/lib/python3.9/site-packages/transformers/models/segformer/modeling_tf_segformer.py", line 547, in serving  *
        output = self.call(inputs)
    File "/home/maxime/anaconda3/envs/tf/lib/python3.9/site-packages/transformers/models/segformer/modeling_tf_segformer.py", line 753, in call  *
        batch_size = shape_list(encoder_hidden_states[-1])[0]

    KeyError: -1
```

### Minimum reproducting code:

```
import os
import tensorflow as tf
from transformers import TFSegformerForSemanticSegmentation


model = TFSegformerForSemanticSegmentation.from_pretrained(
    "nvidia/mit-b0",
    num_labels=2,
    id2label={1:"1", 2:"2"},
    ignore_mismatched_sizes=True,  # Will ensure the segmentation specific components are reinitialized.

)

model.summary(line_length=250)
tf.keras.models.save_model(
    model, os.path.join("/tmp", model.name), include_optimizer=False
)
```

### Expected behavior
The checkpoint should save correctly

### System Info
Ubuntu 20.04, Python 3.9, TF 2.9.1, Nvidia Titan

### Who can help?

@sayakpaul  @NielsRogge 
Have you tried using `model.save_pretrained()`?

Cc: @amyeroberts <|||||>`model.save_pretrained()` indeed works, thanks :)

Shall we close this issue, or is this model also supposed to be compatible with keras saving method? (their callback is widely used, https://www.tensorflow.org/api_docs/python/tf/keras/callbacks/ModelCheckpoint)<|||||>@joihn Glad to hear you were able to save with `save_pretrained` and thanks for responding so quickly @sayakpaul!

I'll defer the issue to our TF gurus @Rocketknight1 @gante regarding compatibility with keras saving. <|||||>I see two possible options for the time being:

* Implement a custom callback to use `save_pretrained()`. Shouldn't differ too much from the `ModelCheckpoint` callback. 
*  You can refer to [this notebook](https://github.com/huggingface/notebooks/blob/main/examples/image_classification-tf.ipynb) that makes use of `PushToHubCallback()` and achieves a similar result as `ModelCheckpoint` barring some differences. <|||||>I did some digging, this issue seems to appears when keras function [save_model] https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model) parmeters is `save_traces=True`.

It's worth noting that hugging face `saved_pretrained(filepath, saved_model=True)` also crash with the same error.

also related to https://github.com/huggingface/transformers/issues/13742
<|||||>Hi @joihn - there are some general difficulties when saving Hugging Face models as SavedModel. This is a general issue with any model where the model and layers are implemented by subclassing in Keras - SavedModel doesn't really have a good way to completely save and load those models (although you can save one or more model traces through SavedModel, this isn't usually what people want unless they're trying to export to TFLite or something!)

Instead, we recommend that users save weights only, and if they want to save the entire model, to use the `save_pretrained` method, which will save the weights along with a config that will make it loadable with the `load_pretrained` method.

Concretely, this means doing the following things:

1) When using `ModelCheckpoint`, set `save_weights_only` to `True`.
2) Replace `model.save` with either `model.save_weights` or `model.save_pretrained`<|||||>Perfect, thanks for the info :) 