Update test.py

test.py

@@ -35,6 +35,7 @@ def main():
         type=str,
     )
     parser.add_argument("--local_rank", type=int, default=0)
+    parser.add_argument('--local-rank', type=int, default=0)
     parser.add_argument(
         "--ckpt",
         help="The path to the checkpoint for test, default is the latest checkpoint.",
@@ -51,6 +52,8 @@ def main():
         nargs=argparse.REMAINDER,
     )
     args = parser.parse_args()
+    if 'LOCAL_RANK' not in os.environ:
+	    os.environ['LOCAL_RANK'] = str(args.local_rank)
 
     num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
     distributed = num_gpus > 1