def to_bert_input(token_idx, null_idx):
    """
    token_idx is a 2D tensor int.

    return token_idx, segment_idx and mask
    """
    segment_idx = token_idx * 0
    mask = token_idx != null_idx
    token_idx = token_idx * mask.long()
    return token_idx, segment_idx, mask
