def test_t5_forward_loss(
    model: T5,
    hf_model: hf_t5.T5ForConditionalGeneration,
    tokenizer: hf_t5.T5Tokenizer,
):
    """
    Make sure loss and generation match the implementation from HF.
    """
    input_ids = tokenizer(
        ["The <extra_id_0> walks in <extra_id_1> park", "The <extra_id_0> barked"],
        return_tensors="pt",
        padding=True,
    ).input_ids
    assert input_ids.tolist() == [
        [37, 32099, 10681, 16, 32098, 2447, 1],
        [37, 32099, 1207, 5100, 1, 0, 0],
    ]

    attention_mask = ~(input_ids == 0)

    labels = tokenizer(
        ["<extra_id_0> cute dog <extra_id_1> the <extra_id_2>", "<extra_id_0> dog"],
        return_tensors="pt",
        padding=True,
    ).input_ids
    assert labels.tolist() == [
        [32099, 5295, 1782, 32098, 8, 32097, 1],
        [32099, 1782, 1, 0, 0, 0, 0],
    ]

    decoder_attention_mask = ~(labels == 0)

    outputs = model(
        input_ids,
        attention_mask=attention_mask,
        labels=labels,
        decoder_attention_mask=decoder_attention_mask,
    )
    hf_outputs = hf_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
        labels=labels,
        decoder_attention_mask=decoder_attention_mask,
    )

    assert outputs.loss == hf_outputs.loss
    assert (outputs.logits == hf_outputs.logits).all()
Make sure loss and generation match the implementation from HF.
tests/modules/transformer/t5_test.py