passes local rank via environment variable (#1784)

src/sparseml/pytorch/torchvision/train.py

@@ -400,7 +400,7 @@ def collate_fn(batch):
     )
 
     _LOGGER.info("Creating model")
-    local_rank = args.local_rank if args.distributed else None
+    local_rank = int(os.environ["LOCAL_RANK"]) if args.distributed else None
     model, arch_key, maybe_dp_device = _create_model(
         arch_key=args.arch_key,
         local_rank=local_rank,
@@ -1256,14 +1256,6 @@ def new_func(*args, **kwargs):
         "Note: Will use ImageNet values if not specified."
     ),
 )
-@click.option(
-    "--local_rank",
-    "--local-rank",
-    type=int,
-    default=None,
-    help="Local rank for distributed training",
-    hidden=True,  # should not be modified by user
-)
 @click.pass_context
 def cli(ctx, **kwargs):
     """