def call(self, inputs, training=False):
    (x, attention_mask, head_mask) = inputs
    output_attn = self.attn([x, attention_mask, head_mask], training=training)
    a = output_attn[0]
    n = self.ln_1((x + a))
    m = self.mlp(n, training=training)
    h = self.ln_2((n + m))
    outputs = ([h] + output_attn[1:])
    return outputs
def call(self, inputs, training=False):
    (x, attention_mask, head_mask) = inputs
    output_attn = self.attn([x, attention_mask, head_mask], training=training)
    a = output_attn[0]
    n = self.ln_1((x + a))
    m = self.mlp(n, training=training)
    h = self.ln_2((n + m))
    outputs = ([h] + output_attn[1:])
    return outputs