Support mixed precision FP16 in TF Segformer | Nan loss
### Description

I'm trying to bring mixed precision training (FP16) support to [TF Segformer](https://github.com/huggingface/transformers/blob/main/src/transformers/models/segformer/modeling_tf_segformer.py) .
I had to do a very small modification to the inital code for supporting Fp16 ([my commit ](https://github.com/huggingface/transformers/commit/41d2bd145111a743a577f15955bac58a74871b33)).

Everytime, the net converge fine for the first 10-15 epoch, and then the loss suddenly goes to Nan.
Any idea ? 

I'm using 
`policy = mixed_precision.Policy("mixed_float16")` ([TF doc here ](https://www.tensorflow.org/guide/mixed_precision))


### System
Tensorflow version 2.8.2
Cuda 11.6
Nvidia titan X

### Who can help?

@sayakpaul @NielsRogge 


### Reproduction

1) Modify 1 line of code as in ([my commit ](https://github.com/huggingface/transformers/commit/41d2bd145111a743a577f15955bac58a74871b33)).
2) Launch a training with mixed precision FP16
3) wait 10~15 epoch

### Expected behavior

The net shoudl train fine with FP16 mixed precision
Thanks! Are you casting the softmaxed outputs or the logits to FP32 before loss computation as recommended [here](https://www.tensorflow.org/guide/mixed_precision#building_the_model)? 

Nit: You didn't provide any training code. I guess it's always better to provide a self-contained Colab Notebook or a notebook that anyone can spin up to debug the issue further. I also acknowledge that Colab doesn't always provide a GPU that has support for Tensor cores but I hope you got the point. <|||||>thanks for the quick answer, 
Even though my commit above doesn't include it, yes I tried casting the logits output to FP32, ans sadly it didn't solve it :/ 

Good point about the small code example for debugging, I will do one today :) <|||||>Then we need to inspect the layer-wise activations and the weight distributions that are probably impacted by the casting. Maybe try using the TensorBoard callback and inspect if that's the case. <|||||>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.