def forward(self, x, y=None):
    x = self.attn1(x)
    if hasattr(self, 'attn2'):
        x = self.attn2(x, y)
    return x
def forward(self, x, y=None):
    x = self.attn1(x)
    if hasattr(self, 'attn2'):
        x = self.attn2(x, y)
    return x