def METHOD_NAME(N, mode, control_variate_update_freq=10, chunk_size=97):
    hi, rotations, training_samples, rho = _setup_measurements(N, mode)
    if mode == "pure":
        ma = nk.models.RBM(alpha=1, param_dtype=complex)
        sa = nk.sampler.MetropolisLocal(hilbert=hi.physical)
        vs = nk.vqs.MCState(sa, ma, n_samples=1000, seed=SEED)
    elif mode == "mixed":
        ma = nk.models.NDM(alpha=1, param_dtype=complex)
        sa = nk.sampler.MetropolisLocal(hilbert=hi)
        vs = nk.vqs.MCMixedState(sa, ma, n_samples=1000, seed=SEED)
    else:
        raise ValueError("Invalid mode")
    op = nk.optimizer.Adam(learning_rate=0.01)
    driver = nkx.QSR(
        (training_samples, rotations),
        training_batch_size=100,
        optimizer=op,
        variational_state=vs,
        control_variate_update_freq=control_variate_update_freq,
        chunk_size=chunk_size,
    )
    return driver, rho