[T5] Fix init in TF and Flax for pretraining (#17294)
* fix init
* Apply suggestions from code review
* fix
* finish
* Update src/transformers/modeling_tf_utils.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
