def test_summary_layer_types(max_depth):
    """Test that the summary displays the layer names correctly."""
    model = UnorderedModel()
    summary = summarize(model, max_depth=max_depth)
    assert summary.layer_types == ["Linear", "Linear", "Linear", "ReLU", "Conv2d"]
Test that the summary displays the layer names correctly.
tests/tests_pytorch/utilities/test_model_summary.py