class_docstr: Wrapper created to ease the training/testing/loading.
  
  Args:
      model (nn.Module): The model to optimize.
      criterion (Callable): A loss function.
      replicate_in_memory (bool): Replicate in memory optional.
class_docstr_tok: ["Wrapper","created","to","ease","the","training\/testing\/loading",".","Args",":","model","(","nn.Module",")",":","The","model","to","optimize",".","criterion","(","Callable",")",":","A","loss","function",".","replicate_in_memory","(","bool",")",":","Replicate","in","memory","optional","."]
docstr: Use the model to predict on a dataset `iterations` time.
  
  Args:
      dataset (Dataset): Dataset to predict on.
      batch_size (int):  Batch size to use during prediction.
      iterations (int): Number of iterations per sample.
      use_cuda (bool): Use CUDA or not.
      workers (int): Number of workers to use.
      collate_fn (Optional[Callable]): The collate function to use.
      half (bool): If True use half precision.
      verbose (bool): If True use tqdm to show progress.
  
  Notes:
      The "batch" is made of `batch_size` * `iterations` samples.
  
  Returns:
      Array [n_samples, n_outputs, ..., n_iterations].
docstr_tok: ["Use","the","model","to","predict","on","a","dataset","`","iterations","`","time",".","Args",":","dataset","(","Dataset",")",":","Dataset","to","predict","on",".","batch_size","(","int",")",":","Batch","size","to","use","during","prediction",".","iterations","(","int",")",":","Number","of","iterations","per","sample",".","use_cuda","(","bool",")",":","Use","CUDA","or","not",".","workers","(","int",")",":","Number","of","workers","to","use",".","collate_fn","(","Optional","[","Callable","]",")",":","The","collate","function","to","use",".","half","(","bool",")",":","If","True","use","half","precision",".","verbose","(","bool",")",":","If","True","use","tqdm","to","show","progress",".","Notes",":","The","``","batch","''","is","made","of","`","batch_size","`","*","`","iterations","`","samples",".","Returns",":","Array","[","n_samples",",","n_outputs",",","...",",","n_iterations","]","."]
code: def predict_on_dataset(
      self,
      dataset: Dataset,
      batch_size: int,
      iterations: int,
      use_cuda: bool,
      workers: int = 4,
      collate_fn: Optional[Callable] = None,
      half=False,
      verbose=True,
  ):
      """
      Use the model to predict on a dataset `iterations` time.
  
      Args:
          dataset (Dataset): Dataset to predict on.
          batch_size (int):  Batch size to use during prediction.
          iterations (int): Number of iterations per sample.
          use_cuda (bool): Use CUDA or not.
          workers (int): Number of workers to use.
          collate_fn (Optional[Callable]): The collate function to use.
          half (bool): If True use half precision.
          verbose (bool): If True use tqdm to show progress.
  
      Notes:
          The "batch" is made of `batch_size` * `iterations` samples.
  
      Returns:
          Array [n_samples, n_outputs, ..., n_iterations].
      """
      preds = list(
          self.predict_on_dataset_generator(
              dataset=dataset,
              batch_size=batch_size,
              iterations=iterations,
              use_cuda=use_cuda,
              workers=workers,
              collate_fn=collate_fn,
              half=half,
              verbose=verbose,
          )
      )
  
      if len(preds) > 0 and not isinstance(preds[0], Sequence):
          # Is an Array or a Tensor
          return np.vstack(preds)
      return [np.vstack(pr) for pr in zip(*preds)]
  
imports: import sys
  from collections import defaultdict
  from collections.abc import Sequence
  from copy import deepcopy
  from typing import Callable, Optional
  import numpy
  import structlog
  import torch
  from torch.optim import Optimizer
  from torch.utils.data import DataLoader
  from torch.utils.data.dataloader import default_collate
  from tqdm import tqdm
  from baal.metrics.mixin import MetricMixin
  from baal.utils.array_utils import stack_in_memory
  from baal.active.dataset.base import Dataset
  from baal.utils.cuda_utils import to_cuda
  from baal.utils.iterutils import map_on_tensor
  from baal.utils.metrics import Loss
  from baal.utils.warnings import raise_warnings_cache_replicated
label_desc: Use image node_id 6 for calling the ModelWrapper obj's underlying member method code with example usage: obj.predict_on_dataset(dataset, batch_size, iterations, use_cuda, workers, collate_fn, half, verbose) and returns: unknown, np