How to conver Pytorch model to Flax

I want to train a mT5 model on more data for a sepcific language, how can I resume training (not from scratch)? I found a training example in huggingface examples, but that is for flax and from scratch. If I change FlaxT5ForConditionalGeneration to T5ForConditionalGenration and load my pytorch mode, then would it work? If no, how can I convert my pytorch to flax so that I can use this code?

I believe you can convert a model from one framework to the other as follows:

from transformers import T5ForConditionalGeneration

model = T5ForConditionalGeneration.from_pretrained(model_name, from_flax=True)