Force minimize_nested_blockmodel_dl to find a specific number of blocks

I am aware that for minimize_blockmodel_dl, one can specify B_min and B_max to a specific number, i.e.,

gt.minimize_blockmodel_dl(g, state=gt.BlockState, multilevel_mcmc_args=dict(B_min=n_blocks, B_max=n_blocks))

will force SBM to find n_blocks blocks.

However, will doing the same thing for minimize_nested_blockmodel_dl work as intended? Because the source code for that is

def minimize_nested_blockmodel_dl(
    g,
    init_bs=None,
    state=gt.NestedBlockState,
    state_args={},
    multilevel_mcmc_args={},
):
    state = state(g, bs=init_bs, **state_args)

    args = dict(niter=1, beta=np.inf)
    args.update(multilevel_mcmc_args)

    l = 0
    while l >= 0:
        ret = state.multilevel_mcmc_sweep(ls=[l], **args)

        if args.get("verbose", False):
            print(l, ret, state)

        if abs(ret[0]) < 1e-8:
            l -= 1
        else:
            l = min((l + 1, len(state.levels) - 1))

    return state

which looks to me like it will force that number of blocks for all layers.

If my interpretation above is correct, will it be reasonable to write a custom function modifying the above to, say,

def minimize_nested_blockmodel_dl(
    g,
    init_bs=None,
    state=gt.NestedBlockState,
    state_args={},
    multilevel_mcmc_args={},
    nblocks=None,
):
    state = state(g, bs=init_bs, **state_args)

    args = dict(niter=1, beta=np.inf)
    args.update(multilevel_mcmc_args)

    l = 0
    while l >= 0:
        args_copy = args.copy()
        if l == 0 and nblocks is not None:
            args_copy.update(B_min=nblocks, B_max=nblocks)
        ret = state.multilevel_mcmc_sweep(ls=[l], **args_copy)

        if args.get("verbose", False):
            print(l, ret, state)

        if abs(ret[0]) < 1e-8:
            l -= 1
        else:
            l = min((l + 1, len(state.levels) - 1))

    return state

to achieve what I want?

Thank you very much!

It will work exactly as intended for the nested version as well. No need at all to modify the code.