No attribute '_mp_fn' when fine-tuning mbart for en-ro translation task using TPU

I followed the TPU example in the examples folder and found xla_spawn.py calls xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
and fine-tune.py does not have the “_mp_fn” found in some training scripts.
I get

Traceback (most recent call last):
  File "examples/xla_spawn.py", line 72, in <module>
    main()
  File "examples/xla_spawn.py", line 68, in main
    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
AttributeError: module 'finetune' has no attribute '_mp_fn'

Tried to fix it by adding the _mp_fn:

#def _mp_fn(index):
    # For xla_spawn (TPUs)
#    pass
    #main()

with and without args main(args) but neither worked.

[sshleifer]

This is now supported by Seq2SeqTrainer which doesn’t use PL. See https://github.com/huggingface/transformers/blob/master/examples/seq2seq/builtin_trainer/finetune_tpu.sh

[abedkhooli]

So, lightning_base.py is not ready for TPU yet.