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

Small optimization for LSTM

parent 8fcdb9b7

fairseq/models/lstm.py

+1 −1
Changes for fairseq/models/lstm.py: 1 added line, 1 removed line.
Original line number Diff line number Diff line
@@ -245,7 +245,7 @@ class AttentionLayer(nn.Module): 
                float('-inf')

            ).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

        x = (attn_scores.unsqueeze(2) * source_hids).sum(dim=0)