fix bug for gen conversion

Author: levindabhi

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)
-    )
+    )