local_rank fix (#1129)

xwang233 · web-flow · commit c1378e6f6224 · 2021-07-16T14:02:26.000-07:00
* local_rank and install cuda version fix
diff --git a/examples/imagenet/main_amp.py b/examples/imagenet/main_amp.py
@@ -81,7 +81,7 @@ def parse():
                         help='Only run 10 iterations for profiling.')
     parser.add_argument('--deterministic', action='store_true')
 
-    parser.add_argument("--local_rank", default=0, type=int)
+    parser.add_argument("--local_rank", default=os.getenv('LOCAL_RANK', 0), type=int)
     parser.add_argument('--sync_bn', action='store_true',
                         help='enabling apex sync BN.')
 
