def test_backward_required(_, strategy, precision, error_expected, setup_method):
    """Test under which strategy and precision configurations the `fabric.backward()` call is required."""
    fabric = Fabric(
        accelerator=("cuda" if strategy == "fsdp" else "cpu"), strategy=strategy, precision=precision, devices=1
    )
    fabric._launched = True
    fabric.strategy.setup_module = lambda module: module

    error_context = (
        pytest.raises(RuntimeError, match=escape("requires you to call `fabric.backward(loss)`"))
        if error_expected
        else nullcontext()
    )
    batch = torch.rand(2, 2)

    # One model
    model1 = nn.Linear(2, 2)
    model1 = getattr(fabric, setup_method)(model1)
    output = model1(batch)
    assert output._backward_hooks is not None
    loss = output.sum()
    with error_context:
        loss.backward()
    loss = model1(batch).sum()
    assert not lightning.fabric.wrappers._in_fabric_backward
    fabric.backward(loss)  # no error
    assert not lightning.fabric.wrappers._in_fabric_backward

    # Two models chained
    model2 = torch.nn.Linear(2, 2)
    model2 = getattr(fabric, setup_method)(model2)
    output = model2(model1(batch))
    assert output._backward_hooks is not None
    loss = output.sum()
    with error_context:
        loss.backward()
    loss = model2(model1(batch)).sum()
    fabric.backward(loss)  # no error

    # Two independent models
    loss1 = model1(batch).sum()
    loss2 = model2(batch).sum()
    with error_context:
        loss1.backward()
    with error_context:
        loss2.backward()
    loss1 = model1(batch).sum()
    loss2 = model2(batch).sum()
    fabric.backward(loss1)  # no error
    fabric.backward(loss2)  # no error

    # Model that returns a datastructure of tensors
    class DictReturnModel(nn.Linear):
        def forward(self, x):
            return {
                "loss": super().forward(x).sum(),
                "other": torch.rand(2, 2),  # does not require grad
            }

    model3 = DictReturnModel(2, 2)
    model3 = getattr(fabric, setup_method)(model3)
    output = model3(batch)
    loss = output["loss"]
    other = output["other"]
    assert loss._backward_hooks is not None
    assert other._backward_hooks is None

    with error_context:
        (loss * 2).backward()
    loss = model3(batch)["loss"]
    fabric.backward(loss * 2)  # no error
Test under which strategy and precision configurations the `fabric.backward()` call is required.
tests/tests_fabric/test_fabric.py