def test_max_depth_param(max_depth):
    """Test that only the modules up to the desired depth are shown."""
    model = DeepNestedModel()
    summary = ModelSummary(model, max_depth=max_depth)
    for lname in summary.layer_names:
        if max_depth >= 0:
            assert lname.count(".") < max_depth
Test that only the modules up to the desired depth are shown.
tests/tests_pytorch/utilities/test_model_summary.py