def test_summary_parameter_count(max_depth):
    """Test that the summary counts the number of parameters in every submodule."""
    model = UnorderedModel()
    summary = summarize(model, max_depth=max_depth)
    assert summary.param_nums == [
        model.layer2.weight.numel() + model.layer2.bias.numel(),
        model.combine.weight.numel() + model.combine.bias.numel(),
        model.layer1.weight.numel() + model.layer1.bias.numel(),
        0,  # ReLU
        model.unused.weight.numel() + model.unused.bias.numel(),
    ]
Test that the summary counts the number of parameters in every submodule.
tests/tests_pytorch/utilities/test_model_summary.py