test

mar.py

@@ -309,6 +309,8 @@ class MARBert(nn.Module):
                 mask_to_pred = torch.cat([mask_to_pred, mask_to_pred], dim=0)
 
             # sample token latents for this step
+            print(z.shape)
+            print("-----------")
             z = z[mask_to_pred.nonzero(as_tuple=True)]
             print(z.shape)
             # cfg schedule follow Muse