test

mar.py

@@ -117,7 +117,7 @@ class MARBert(nn.Module):
         print("test")
         print("test")
         print("test")
-        
+
     def initialize_weights(self):
         # parameters
         torch.nn.init.normal_(self.class_emb.weight, std=.02)
@@ -599,7 +599,7 @@ class MAR(nn.Module):
         for step in indices:
             cur_tokens = tokens.clone()
             print(cur_tokens.shape)
-
+            print("+++++++++++++++++++++++")
             # class embedding and CFG
             if labels is not None:
                 class_embedding = self.class_emb(labels)
@@ -635,7 +635,10 @@ class MAR(nn.Module):
                 mask_to_pred = torch.cat([mask_to_pred, mask_to_pred], dim=0)
 
             # sample token latents for this step
+            print(z.shape)
             z = z[mask_to_pred.nonzero(as_tuple=True)]
+            print(z.shape)
+            print("==============================")
             # cfg schedule follow Muse
             if cfg_schedule == "linear":
                 cfg_iter = 1 + (cfg - 1) * (self.seq_len - mask_len[0]) / self.seq_len