def METHOD_NAME(mod, dataset):
    cfg = quantize.current_qconfig()
    chunk_by = cfg.calibrate_chunk_by
    scales = []
    for samples in collect_stats(mod, dataset, chunk_by):
        logging.info("finding threshold with kl for calibration...")
        with mp.Pool() as pool:
            scales += list(pool.map(_find_scale_by_kl, samples))
    def func(_):
        scale = scales[func.scale_idx]
        func.scale_idx += 1
        return scale
    func.scale_idx = 0
    return func