Fix the missing parameter error when running mp_imagenet with torchrun (#5729)

* Fix the missing parameter error when running mp_imagenet with torchrun * made it local rank

Author: vanbasten23

test/test_train_mp_imagenet.py

@@ -85,6 +85,7 @@
 import torch_xla.distributed.parallel_loader as pl
 import torch_xla.debug.profiler as xp
 import torch_xla.utils.utils as xu
+import torch_xla.core.xla_env_vars as xenv
 import torch_xla.core.xla_model as xm
 import torch_xla.distributed.xla_multiprocessing as xmp
 import torch_xla.test.test_utils as test_utils
@@ -375,6 +376,6 @@ def _mp_fn(index, flags):
 
 if __name__ == '__main__':
  if dist.is_torchelastic_launched():
- _mp_fn(FLAGS)
+ _mp_fn(xu.getenv_as(xenv.LOCAL_RANK, int), FLAGS)
  else:
  xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=FLAGS.num_cores)