From da129e0304b5a129c3793055b03647703d2b43f1 Mon Sep 17 00:00:00 2001
From: lucidrains <lucidrains@gmail.com>
Date: Tue, 24 Oct 2023 08:21:06 -0700
Subject: [PATCH] ready for training

---
 magvit2_pytorch/magvit2_pytorch.py | 13 +++++++++++++
 setup.py                           |  1 +
 2 files changed, 14 insertions(+)

diff --git a/magvit2_pytorch/magvit2_pytorch.py b/magvit2_pytorch/magvit2_pytorch.py
index f7d81be..bb15186 100644
--- a/magvit2_pytorch/magvit2_pytorch.py
+++ b/magvit2_pytorch/magvit2_pytorch.py
@@ -1036,6 +1036,19 @@ def __init__(
 
         self.has_gan = use_gan and adversarial_loss_weight > 0.
 
+    def parameters(self):
+        return [
+            *self.conv_in.parameters(),
+            *self.conv_out.parameters(),
+            *self.encoder_layers.parameters(),
+            *self.decoder_layers.parameters(),
+            *self.encoder_cond_in.parameters(),
+            *self.decoder_cond_in.parameters(),
+        ]
+
+    def discr_parameters(self):
+        return self.discr.parameters()
+
     def copy_for_eval(self):
         device = next(self.parameters()).device
         vae_copy = copy.deepcopy(self.cpu())
diff --git a/setup.py b/setup.py
index 86bb826..1ebfcef 100644
--- a/setup.py
+++ b/setup.py
@@ -21,6 +21,7 @@
     'accelerate',
     'beartype',
     'einops>=0.7.0',
+    'ema-pytorch',
     'kornia',
     'vector-quantize-pytorch>=1.9.18',
     'torch',