def process_batch(batch):
    """Process batch and produce inputs for the model."""

    loss_mask = batch['pad_mask'].long().contiguous().byte()
    tokens_ = batch['text'].long().contiguous()
    labels = tokens_[:, 1:].contiguous()
    tokens = tokens_[:, :-1].contiguous()

    return tokens, labels, loss_mask