Added option to use the ema_model argument when initializing the EMA …

…from the normal VAE, this way we can pass an already initialized EMA VAE to the trainer.
ZeroCool940711 · Sep 5, 2023 · a1a8de1 · a1a8de1
1 parent 09e5f92
commit a1a8de1
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/muse_maskgit_pytorch/trainers/vqvae_trainers.py b/muse_maskgit_pytorch/trainers/vqvae_trainers.py
@@ -55,6 +55,7 @@ def __init__(
         lr_warmup_steps=500,
         discr_max_grad_norm=None,
         use_ema=True,
+        ema_vae=None,
         ema_beta=0.995,
         ema_update_after_step=0,
         ema_update_every=1,
@@ -155,6 +156,7 @@ def __init__(
         if use_ema:
             self.ema_model = EMA(
                 vae,
+                ema_model=ema_vae,
                 update_after_step=ema_update_after_step,
                 update_every=ema_update_every,
             )