def METHOD_NAME(model: torch.nn.Module,
                        dataloader: DataLoader,
                        num_batches: int = DEFAULT_NUM_BATCHES,
                        forward_fn: Callable[[torch.nn.Module, Any], Any] = None) -> Handle:
    """
    Reestimate BatchNorm statistics (running mean and var).
    :param model: Model to reestimate the BN stats.
    :param dataloader: Training dataset.
    :param num_batches: The number of batches to be used for reestimation.
    :param forward_fn: Optional adapter function that performs forward pass
                       given a model and a input batch yielded from the data loader.
    :returns: Handle that undos the effect of BN reestimation upon handle.remove().
    """
    forward_fn = forward_fn or (lambda model, data: model(data))
    bn_modules = tuple(_get_active_bn_modules(model))
    # Set all the layers to eval mode except batchnorm layers
    with in_eval_mode(model), in_train_mode(bn_modules), torch.no_grad():
        with _for_each_module(bn_modules, action=_reset_momentum):
            handle = _for_each_module(bn_modules, action=_reset_bn_stats)
            try:
                # Batchnorm statistics accumulation buffer
                buffer = {
                    bn: {"sum_mean": torch.zeros_like(bn.running_mean),
                         "sum_var":  torch.zeros_like(bn.running_var)}
                    for bn in bn_modules
                }
                num_batches = min(len(dataloader), num_batches)
                dataloader_slice = itertools.islice(dataloader, num_batches)
                for data in tqdm(dataloader_slice,
                                 total=num_batches,
                                 desc="batchnorm reestimation"):
                    forward_fn(model, data)
                    for bn in bn_modules:
                        buffer[bn]["sum_mean"] += bn.running_mean
                        buffer[bn]["sum_var"] += bn.running_var
                for bn in bn_modules:
                    sum_mean = buffer[bn]["sum_mean"]
                    sum_var = buffer[bn]["sum_var"]
                    # Override BN stats with the reestimated stats.
                    bn.running_mean.copy_(sum_mean / min(len(dataloader), num_batches))
                    bn.running_var.copy_(sum_var / min(len(dataloader), num_batches))
                return handle
            except:
                handle.remove()
                raise