(self, x, padding_mask=None)
return h
Input has shape [seq length, batch]
Input has shape [seq length, batch]
Input has shape [ seq length batch ]
def forward(self, x, padding_mask=None):
        """ Input has shape [seq length, batch] """
        positions = torch.arange(len(x), device=x.device).unsqueeze(-1)
        h = self.tokens_embeddings(x)
        h = h + self.position_embeddings(positions).expand_as(h)
        h = self.dropout(h)

        attn_mask = None
        if self.causal:
            attn_mask = torch.full((len(x), len(x)), -float('Inf'), device=h.device, dtype=h.dtype)
            attn_mask = torch.triu(attn_mask, diagonal=1)
        for layer_norm_1, attention, layer_norm_2, feed_forward in zip(self.layer_norms_1, self.attentions,
                                                                       self.layer_norms_2, self.feed_forwards):
            h = layer_norm_1(h)
            x, _ = attention(h, h, h, attn_mask=attn_mask, need_weights=False, key_padding_mask=padding_mask)
            x = self.dropout(x)
            h = x + h

            h = layer_norm_2(h)
            x = feed_forward(h)
            x = self.dropout(x)
            h = x + h
        return h
def forward ( self , x , padding_mask = None ) : positions = torch . arange ( len ( x ) , device = x . device ) . unsqueeze ( - 1 ) h = self . tokens_embeddings ( x ) h = h + self . position_embeddings ( positions ) . expand_as ( h ) h = self . dropout ( h ) attn_mask = None if self . causal : attn_mask = torch . full ( ( len ( x ) , len ( x ) ) , - float ( 'Inf' ) , device = h . device , dtype = h . dtype ) attn_mask = torch . triu ( attn_mask , diagonal = 1 ) for layer_norm_1 , attention , layer_norm_2 , feed_forward in zip ( self . layer_norms_1 , self . attentions , self . layer_norms_2 , self . feed_forwards ) : h = layer_norm_1 ( h ) x , _ = attention ( h , h , h , attn_mask = attn_mask , need_weights = False , key_padding_mask = padding_mask ) x = self . dropout ( x ) h = x + h h = layer_norm_2 ( h ) x = feed_forward ( h ) x = self . dropout ( x ) h = x + h return h