def METHOD_NAME(enc_input, attn_mask):
    encoder_layer = paddle.nn.TransformerEncoderLayer(128, 2, 512)
    encoder = paddle.nn.TransformerEncoder(encoder_layer, 2)
    encoder_output = encoder(enc_input, attn_mask)
    first_token = encoder_output[:, 0]
    bias = paddle.full(shape=[1, 128], fill_value=1e-6)
    linear = paddle.nn.Linear(128, 2)
    logits = linear(first_token + bias)
    return logits