reshape logits

config.json

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

modeling_flexbert.py

@@ -1702,12 +1702,9 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 shift_labels.view(-1)
             )
 
-        # if self.unpad_embeddings:
-        #     # reshape to batch size
-        #     logits = logits.view(-1, self.vocab_size)
-        #     # NOTE: error from here above
+        if self.unpad_embeddings:
+            logits = logits.view(batch_size, -1, self.vocab_size)
 
-        breakpoint()
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len-1)[0]