return past hidden states when `output_hidden_states` provided

modeling_phi.py

@@ -947,6 +947,7 @@ class PhiForCausalLM(PhiPreTrainedModel):
         input_ids: torch.LongTensor,
         past_key_values: Optional[Union[torch.FloatTensor, InferenceParams]] = None,
         attention_mask: Optional[torch.BoolTensor] = None,
+        output_hidden_states: Optional[bool] = None,
         labels: Optional[torch.LongTensor] = None,
         **kwargs,
     ) -> CausalLMOutputWithPast:
@@ -957,4 +958,4 @@ class PhiForCausalLM(PhiPreTrainedModel):
         if labels is not None:
             loss = self.loss(lm_logits, labels)
 
-        return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values)
+        return CausalLMOutputWithPast(loss=loss, logits=lm_logits, past_key_values=past_key_values, hidden_states=hidden_states if output_hidden_states else None)