diff --git a/tools/test.py b/tools/test.py index ada5e607d2ca8bc02681c2d999be913e8328d5bd..af950aac2918e81af3d1ab200c0d74d492828ae5 100644 --- a/tools/test.py +++ b/tools/test.py @@ -1,4 +1,5 @@ import argparse +import os import os.path as osp import shutil import tempfile @@ -119,6 +120,8 @@ def parse_args(): help='job launcher') parser.add_argument('--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) return args diff --git a/tools/train.py b/tools/train.py index b8f21d11fa687ffd4e9c5aaf8b5d46b5142c826e..d8bb9dc0f416f398b41ef8af7dd7c3c9743a0838 100644 --- a/tools/train.py +++ b/tools/train.py @@ -1,6 +1,7 @@ from __future__ import division import argparse +import os from mmcv import Config from mmdet import __version__ @@ -35,6 +36,8 @@ def parse_args(): help='job launcher') parser.add_argument('--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) return args