Skip to content

Commit

Permalink
[Feats]: allow setting local-rank for pytorch 2.0 (#2387)
Browse files Browse the repository at this point in the history
  • Loading branch information
sunjiahao1999 authored Mar 28, 2023
1 parent c9fd5a6 commit 20987e5
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion tools/train.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,7 +50,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)
Expand Down

0 comments on commit 20987e5

Please sign in to comment.