fix: D_global optimization

jolibrain · Dec 16, 2021 · 284d8ae · 284d8ae
1 parent bb366b7
commit 284d8ae
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/models/cut_model.py b/models/cut_model.py
@@ -147,7 +147,10 @@ def __init__(self, opt,rank):
             self.group_G = NetworkGroup(networks_to_optimize=["G","F"],forward_functions=["forward"],backward_functions=["compute_G_loss"],loss_names_list=["loss_names_G"],optimizer=["optimizer_G"],loss_backward=["loss_G"])
             self.networks_groups.append(self.group_G)
 
-            self.group_D = NetworkGroup(networks_to_optimize=["D"],forward_functions=None,backward_functions=["compute_D_loss"],loss_names_list=["loss_names_D"],optimizer=["optimizer_D"],loss_backward=["loss_D_tot"])
+            D_to_optimize = ["D"]
+            if opt.netD_global != "none":
+                D_to_optimize.append("D_global")
+            self.group_D = NetworkGroup(networks_to_optimize=D_to_optimize,forward_functions=None,backward_functions=["compute_D_loss"],loss_names_list=["loss_names_D"],optimizer=["optimizer_D"],loss_backward=["loss_D_tot"])
             self.networks_groups.append(self.group_D)