Skip to content

GitLab

Commit f607d9e8 authored by Myle Ott's avatar Myle Ott
Browse files

Small optimization for LSTM

parent 8fcdb9b7
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
fairseq/models/lstm.py
View file @ f607d9e8
...@@ -245,7 +245,7 @@ class AttentionLayer(nn.Module): ...@@ -245,7 +245,7 @@ class AttentionLayer(nn.Module):
float('-inf') float('-inf')
).type_as(attn_scores) # FP16 support: cast to float and back ).type_as(attn_scores) # FP16 support: cast to float and back
attn_scores = F.softmax(attn_scores.t(), dim=1).t() # srclen x bsz attn_scores = F.softmax(attn_scores, dim=0) # srclen x bsz
# sum weighted sources # sum weighted sources
x = (attn_scores.unsqueeze(2) * source_hids).sum(dim=0) x = (attn_scores.unsqueeze(2) * source_hids).sum(dim=0)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment