Archives/petals
2
0
Fork 0
mirror of https://github.com/bigscience-workshop/petals synced 2024-10-31 09:20:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Don't cast logits to float32 on GPU

Browse Source
This commit is contained in:
Aleksandr Borzunov 2022-11-29 08:58:58 +00:00
parent 6e7565e41e
commit 86d08bf515
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/bloom/model.py
View File

@ -449,7 +449,7 @@ class LMHead(nn.Module):
else:
# Switch dtype in case word_embeddings are fp16/bf16
hidden_states = hidden_states.to(word_embeddings.dtype)
lm_logits = F.linear(hidden_states, word_embeddings).float()
lm_logits = F.linear(hidden_states, word_embeddings)
return lm_logits
def chunked_forward(self, hidden_states):