facepalm, forgot to include quantizer parameters. recon loss now looks normal

Author: lucidrains

magvit2_pytorch/magvit2_pytorch.py

@@ -1353,6 +1353,7 @@ def parameters(self):
             *self.decoder_layers.parameters(),
             *self.encoder_cond_in.parameters(),
             *self.decoder_cond_in.parameters(),
+            *self.quantizers.parameters()
         ]
 
     def discr_parameters(self):

magvit2_pytorch/trainer.py

@@ -253,7 +253,7 @@ def train_step(self, dl_iter):
 
             self.accelerator.backward(loss / self.grad_accum_every)
 
-        self.print(f'loss: {loss.item():.3f}')
+        self.print(f'recon loss: {loss_breakdown.recon_loss.item():.3f}')
 
         if exists(self.max_grad_norm):
             self.accelerator.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
@@ -336,7 +336,7 @@ def valid_step(
             valid_videos.append(valid_video)
             recon_videos.append(recon_video)
 
-        self.print(f'validation loss {recon_loss:.3f}')
+        self.print(f'validation recon loss {recon_loss:.3f}')
 
         if not save_recons:
             return

magvit2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.0.70'
+__version__ = '0.1.0'