From a5b60996e139ac00e27fab5e7b4601c023eb8d3b Mon Sep 17 00:00:00 2001 From: Jose <34888496+Jerry-Master@users.noreply.github.com> Date: Thu, 3 Aug 2023 14:39:39 +0200 Subject: [PATCH] Fixed distributed training bug --- styleavatar/train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/styleavatar/train.py b/styleavatar/train.py index d58fae1..8bd743e 100644 --- a/styleavatar/train.py +++ b/styleavatar/train.py @@ -305,7 +305,7 @@ def train(args, loader, back_generator, face_generator, image_generator, discrim parser.add_argument("--ckpt", type=str, default=None, help="path to the checkpoints to resume training") parser.add_argument("--lr", type=float, default=0.002, help="learning rate") parser.add_argument("--channel_multiplier", type=int, default=2, help="channel multiplier factor for the model. config-f = 2, else = 1") - parser.add_argument("--local_rank", type=int, default=0, help="local rank for distributed training") + parser.add_argument("--local-rank", type=int, default=0, help="local rank for distributed training") args = parser.parse_args() transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)])