Merge pull request Sygil-Dev#21 from bzantium/hotfix/when-not-use-ema

hotfix: add use_ema condition for extracting ema_vae in train_step
ZeroCool940711 · Feb 5, 2023 · d301298 · d301298
2 parents 932087b + 2585eef
commit d301298
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/muse_maskgit_pytorch/trainers.py b/muse_maskgit_pytorch/trainers.py
@@ -275,7 +275,8 @@ def train_step(self):
 
         self.vae.train()
         discr = self.vae.module.discr if self.is_distributed else self.vae.discr
-        ema_vae = self.ema_vae.module if self.is_distributed else self.ema_vae
+        if self.use_ema:
+            ema_vae = self.ema_vae.module if self.is_distributed else self.ema_vae
 
         # logs