address lucidrains/vector-quantize-pytorch#112

lucidrains · May 1, 2024 · b17e1ae · b17e1ae
1 parent 9f49074
commit b17e1ae
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/magvit2_pytorch/trainer.py b/magvit2_pytorch/trainer.py
@@ -180,6 +180,8 @@ def __init__(
 
         self.apply_gradient_penalty_every = apply_gradient_penalty_every
 
+        self.has_multiscale_discrs = self.model.has_multiscale_discrs
+
         # prepare for maybe distributed
 
         (
@@ -200,7 +202,6 @@ def __init__(
 
         # multiscale discr losses
 
-        self.has_multiscale_discrs = self.model.has_multiscale_discrs
         self.multiscale_discr_optimizers = []
 
         for ind, discr in enumerate(self.model.multiscale_discrs):

diff --git a/magvit2_pytorch/version.py b/magvit2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.4.0'
+__version__ = '0.4.1'