def METHOD_NAME(output: dict[str, Tensor], expected: dict[str, Tensor]) -> None:
    for key in expected:
        err = f"output[{key}] != expected[{key}]"
        equal = torch.allclose(output[key], expected[key])
        assert equal, err