[Flax] Refactor MLM  (#12013)
[Flax] Refactor MLM  (#12013)

* fix_torch_device_generate_test

* remove @

* finish refactor

Co-authored-by: Patrick von Platen <cbb7353e6d953ef360baf960c122346276c6e320@huggingface.co>
huggingface/transformers,huggingface/transformers,huggingface/transformers,huggingface/transformers