From ed5a653ecfd81c4d74ce71bc2fc44863b7dcf7c2 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 30 Jan 2024 04:16:39 +1100 Subject: [PATCH] Update mistral.py --- unsloth/models/mistral.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 224bfc09c9..08ee0fae93 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -210,7 +210,12 @@ def MistralForCausalLM_fast_forward( ) hidden_states = outputs[0] - logits = self.lm_head(hidden_states) + if hidden_states.shape[0] == 1: + logits = torch.mv(self.lm_head.weight, hidden_states.ravel()) + logits = logits.unsqueeze(0).unsqueeze(0) + else: + logits = self.lm_head(hidden_states) + pass loss = None if labels is not None: