def test_model_size(max_depth):
    """Test model size is calculated correctly."""
    model = PreCalculatedModel()
    summary = summarize(model, max_depth=max_depth)
    assert model.pre_calculated_model_size == summary.model_size
Test model size is calculated correctly.
tests/tests_pytorch/utilities/test_model_summary.py