def METHOD_NAME(model_fn, ref_output_shapes):
    config = get_basic_pruning_config(8)
    input_shape = [1, 8, 8, 1]
    model = model_fn(input_shape)
    model.compile()
    _, compression_ctrl = create_compressed_model_and_algo_for_test(model, config)
    # pylint:disable=protected-access
    output_shapes = collect_output_shapes(compression_ctrl.model, compression_ctrl._original_graph)
    assert output_shapes == ref_output_shapes