[Fix] Fix local-rank argument in test(#2469)

* fix local rank in test 2 * fix test.py
open-mmlab · Apr 26, 2023 · 74bbc07 · 74bbc07
1 parent 37055ca
commit 74bbc07
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tools/test.py b/tools/test.py
@@ -57,7 +57,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/test.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)