ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update mistral.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-30 04:16:39 +11:00
commit ed5a653ecf
1 changed files with 6 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/models/mistral.py
View file

@ -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: