diff --git a/tools/test.py b/tools/test.py index 64da2bc261..058fdfc864 100644 --- a/tools/test.py +++ b/tools/test.py @@ -47,7 +47,10 @@ def parse_args(): help='job launcher') parser.add_argument( '--tta', action='store_true', help='Test time augmentation') - parser.add_argument('--local_rank', type=int, default=0) + # When using PyTorch version >= 2.0.0, the `torch.distributed.launch` + # will pass the `--local-rank` parameter to `tools/train.py` instead + # of `--local_rank`. + parser.add_argument('--local_rank', '--local-rank', type=int, default=0) args = parser.parse_args() if 'LOCAL_RANK' not in os.environ: os.environ['LOCAL_RANK'] = str(args.local_rank) diff --git a/tools/train.py b/tools/train.py index 1721306664..10fdaa1874 100644 --- a/tools/train.py +++ b/tools/train.py @@ -40,7 +40,10 @@ def parse_args(): choices=['none', 'pytorch', 'slurm', 'mpi'], default='none', help='job launcher') - parser.add_argument('--local_rank', type=int, default=0) + # When using PyTorch version >= 2.0.0, the `torch.distributed.launch` + # will pass the `--local-rank` parameter to `tools/train.py` instead + # of `--local_rank`. + parser.add_argument('--local_rank', '--local-rank', type=int, default=0) args = parser.parse_args() if 'LOCAL_RANK' not in os.environ: os.environ['LOCAL_RANK'] = str(args.local_rank)