From 74bbc075fa79356abcd8ed773dae31d082124192 Mon Sep 17 00:00:00 2001
From: Sun Jiahao <72679458+sunjiahao1999@users.noreply.github.com>
Date: Wed, 26 Apr 2023 09:24:14 +0800
Subject: [PATCH] [Fix] Fix local-rank argument in test(#2469)

* fix local rank in test 2

* fix test.py
---
 tools/test.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tools/test.py b/tools/test.py
index f1e125f26d..4afc25597b 100644
--- 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)