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: Get the model's prediction on a batch.
  
  Args:
      data (Tensor): The model input.
      iterations (int): Number of prediction to perform.
      cuda (bool): Use CUDA or not.
  
  Returns:
      Tensor, the loss computed from the criterion.
              shape = {batch_size, nclass, n_iteration}.
  
  Raises:
      Raises RuntimeError if CUDA rans out of memory during data replication.
docstr_tok: ["Get","the","model","'s","prediction","on","a","batch",".","Args",":","data","(","Tensor",")",":","The","model","input",".","iterations","(","int",")",":","Number","of","prediction","to","perform",".","cuda","(","bool",")",":","Use","CUDA","or","not",".","Returns",":","Tensor",",","the","loss","computed","from","the","criterion",".","shape","=","{","batch_size",",","nclass",",","n_iteration","}",".","Raises",":","Raises","RuntimeError","if","CUDA","rans","out","of","memory","during","data","replication","."]
code: def predict_on_batch(self, data, iterations=1, cuda=False):
      """
      Get the model's prediction on a batch.
  
      Args:
          data (Tensor): The model input.
          iterations (int): Number of prediction to perform.
          cuda (bool): Use CUDA or not.
  
      Returns:
          Tensor, the loss computed from the criterion.
                  shape = {batch_size, nclass, n_iteration}.
  
      Raises:
          Raises RuntimeError if CUDA rans out of memory during data replication.
      """
      with torch.no_grad():
          if cuda:
              data = to_cuda(data)
          if self.replicate_in_memory:
              data = map_on_tensor(
                  lambda d: stack_in_memory(d, iterations), data
              )
              try:
                  out = self.model(data)
              except RuntimeError as e:
                  raise RuntimeError(
                      """CUDA ran out of memory while BaaL tried to replicate data. See the exception above.
                      Use `replicate_in_memory=False` in order to reduce the memory requirements.
                      Note that there will be some speed trade-offs"""
                  ) from e
              out = map_on_tensor(
                  lambda o: o.view([iterations, -1, *o.size()[1:]]), out
              )
              out = map_on_tensor(
                  lambda o: o.permute(
                      1, 2, *range(3, o.ndimension()), 0
                  ),
                  out,
              )
          else:
              out = [self.model(data) for _ in range(iterations)]
              out = _stack_preds(out)
          return out
  
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 9 for calling the ModelWrapper obj's underlying member method code with example usage: obj.predict_on_batch(data, iterations, cuda) and returns: out