fix

config.json

@@ -69,7 +69,7 @@
   "num_attention_heads": 12,
   "num_hidden_layers": 22,
   "num_initial_layers": 1,
-  "pad_logits": false,
+  "pad_logits": true,
   "pad_token_id": 0,
   "padding": "unpadded",
   "pooling_type": "cls",

modeling_flexbert.py

@@ -1701,9 +1701,10 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels.view(-1)
             )
 
-        if self.unpad_embeddings:
-            # reshape to batch size
-            logits = logits.view(-1, self.vocab_size)
+        # if self.unpad_embeddings:
+        #     # reshape to batch size
+        #     logits = logits.view(-1, self.vocab_size)
+        #     # NOTE: error from here above
 
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")