print

modeling_flexbert.py

@@ -1707,8 +1707,9 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         #     # NOTE: error from here above
 
         if self.pad_logits:
-            # print(f"Padding logits: {logits.shape}")
+            print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
+            print(f"New logits: {new_logits.shape}")
             # print(new_logits.shape)
             # if new_logits.dim() == 2:
             #     new_logits = new_logits.unsqueeze(0)