test

mar.py

@@ -310,6 +310,7 @@ class MARBert(nn.Module):
 
             # sample token latents for this step
             z = z[mask_to_pred.nonzero(as_tuple=True)]
+            print(z.shape)
             # cfg schedule follow Muse
             if cfg_schedule == "linear":
                 cfg_iter = 1 + (cfg - 1) * (self.seq_len - mask_len[0]) / self.seq_len