deepseekmirror/DeepSeek-V3
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepSeek-V3.git synced 2025-05-02 16:38:57 -04:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Update model.py

Browse Source 
Enabling mixed precision training to reduce memory usage and potentially speed up training.
This commit is contained in:
Pedro Dessanti 2025-01-27 08:58:59 -03:00 committed by GitHub
parent b5d872ead0
commit 2bf4595d13
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
inference/model.py
View File

@ -777,6 +777,7 @@ class Transformer(nn.Module):
Returns: Returns:
torch.Tensor: Logits tensor of shape (batch_size, vocab_size). torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
""" """
with autocast():
seqlen = tokens.size(1) seqlen = tokens.size(1)
h = self.embed(tokens) h = self.embed(tokens)
freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen] freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]