Firehawk52-backup/llama
28
0
Fork 0
mirror of https://github.com/meta-llama/llama.git synced 2026-01-15 08:22:55 -03:00
Code Issues Actions Packages Projects Releases Wiki Activity

update the code to use the module's __call__

Browse Source
This commit is contained in:
wangzhihong
2024-03-21 10:09:34 +08:00
parent 52afd48b06
commit 1e8375848d
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
llama/model.py
View File

@@ -403,10 +403,10 @@ class TransformerBlock(nn.Module):
torch.Tensor: Output tensor after applying attention and feedforward layers.
"""
h = x + self.attention.forward(
h = x + self.attention(
self.attention_norm(x), start_pos, freqs_cis, mask
)
out = h + self.feed_forward.forward(self.ffn_norm(h))
out = h + self.feed_forward(self.ffn_norm(h))
return out