make sure to actually train the vision aided discriminator

Author: lucidrains

gigagan_pytorch/gigagan_pytorch.py

@@ -2135,10 +2135,14 @@ def train_discriminator_step(
         all_real_images = []
 
         self.G.train()
-        self.D.train()
 
+        self.D.train()
         self.D_opt.zero_grad()
 
+        if self.need_vision_aided_discriminator:
+            self.VD.train()
+            self.VD_opt.zero_grad()
+
         for _ in range(grad_accum_every):
 
             if self.unconditional:
@@ -2344,6 +2348,9 @@ def train_discriminator_step(
 
         self.D_opt.step()
 
+        if self.need_vision_aided_discriminator:
+            self.VD_opt.step()
+
         return TrainDiscrLosses(
             total_divergence,
             total_multiscale_divergence,

gigagan_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.2.12'
+__version__ = '0.2.14'