Merge pull request rosinality#182 from levindabhi/master

fix bug for gen conversion
yeruiqian · Jan 27, 2021 · 13d3a39 · 13d3a39
2 parents c5ba53a + cc57282
commit 13d3a39
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/convert_weight.py b/convert_weight.py
@@ -256,7 +256,7 @@ def fill_statedict(state_dict, vars, size, n_mlp):
     if args.gen:
         g_train = Generator(size, 512, n_mlp, channel_multiplier=args.channel_multiplier)
         g_train_state = g_train.state_dict()
-        g_train_state = fill_statedict(g_train_state, generator.vars, size)
+        g_train_state = fill_statedict(g_train_state, generator.vars, size, n_mlp)
         ckpt["g"] = g_train_state
 
     if args.disc:
@@ -298,4 +298,4 @@ def fill_statedict(state_dict, vars, size, n_mlp):
 
     utils.save_image(
         img_concat, name + ".png", nrow=n_sample, normalize=True, range=(-1, 1)
-    )
+    )