A teeny bug on Multi-GPU training

[EricZhang1412]

There are a small bug in train.py. When I apply the proposed model for multi-gpu training, it got unrecognized error: --local-rank. It seems that the linux system defaultly transmit the config "local-rank" instead of "local_rank" (see the different between '-' and '_'). As soon as i add
parser.add_argument("--local_rank", "--local-rank", default=0, type=int) to parser the two arguments into one.
it went well.

[jkhu29]

Actually, this is due to different pytorch versions. After PyTorch 2.0, pytorch recommends to use "local-rank". Please see [https://pytorch.org/docs/stable/elastic/run.html#transitioning-from-torch-distributed-launch-to-torchrun] 2.