def forward(self, x):
    h = x
    x = self.attention_norm(x)
    (x, weights) = self.attn(x)
    x = (x + h)
    h = x
    x = self.ffn_norm(x)
    x = self.ffn(x)
    x = (x + h)
    return (x, weights)
def forward(self, x):
    h = x
    x = self.attention_norm(x)
    (x, weights) = self.attn(x)
    x = (x + h)
    h = x
    x = self.ffn_norm(x)
    x = self.ffn(x)
    x = (x + h)
    return (x, weights)