logits

modeling_flexbert.py

@@ -1702,7 +1702,8 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
         if self.pad_logits:
             # print(f"Padding logits: {logits.shape}")
             new_logits = self.pad_inputs(logits, indices, batch_size, seq_len)[0]
-            if len(new_logits.shape) == 2:
+            # print(new_logits.shape)
+            if new_logits.dim() == 1:
                 new_logits = new_logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,
@@ -1711,8 +1712,8 @@ class FlexBertForCausalLM(FlexBertPreTrainedModel):
                 attentions=None,
             )
         else:
-            # print(f"Non-padding logits: {logits.shape}")
-            if len(logits.shape) == 2:
+            print(f"Non-padding logits: {logits.shape}")
+            if logits.dim() == 1:
                 logits = logits.unsqueeze(0)
             return CausalLMOutput(
                 loss=loss,