def maybe_layer_norm(self, layer_norm, x, before=False, after=False):
    assert (before ^ after)
    if (after ^ self.normalize_before):
        return layer_norm(x)
    else:
        return x
def maybe_layer_norm(self, layer_norm, x, before=False, after=False):
    assert (before ^ after)
    if (after ^ self.normalize_before):
        return layer_norm(x)
    else:
        return x