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: Train for `epoch` epochs on a Dataset `dataset.
  
  Args:
      dataset (Dataset): Pytorch Dataset to be trained on.
      optimizer (optim.Optimizer): Optimizer to use.
      batch_size (int): The batch size used in the DataLoader.
      epoch (int): Number of epoch to train for.
      use_cuda (bool): Use cuda or not.
      workers (int): Number of workers for the multiprocessing.
      collate_fn (Optional[Callable]): The collate function to use.
      regularizer (Optional[Callable]): The loss regularization for training.
  
  Returns:
      The training history.
docstr_tok: ["Train","for","`","epoch","`","epochs","on","a","Dataset","`","dataset",".","Args",":","dataset","(","Dataset",")",":","Pytorch","Dataset","to","be","trained","on",".","optimizer","(","optim.Optimizer",")",":","Optimizer","to","use",".","batch_size","(","int",")",":","The","batch","size","used","in","the","DataLoader",".","epoch","(","int",")",":","Number","of","epoch","to","train","for",".","use_cuda","(","bool",")",":","Use","cuda","or","not",".","workers","(","int",")",":","Number","of","workers","for","the","multiprocessing",".","collate_fn","(","Optional","[","Callable","]",")",":","The","collate","function","to","use",".","regularizer","(","Optional","[","Callable","]",")",":","The","loss","regularization","for","training",".","Returns",":","The","training","history","."]
code: def train_on_dataset(
      self,
      dataset,
      optimizer,
      batch_size,
      epoch,
      use_cuda,
      workers=4,
      collate_fn: Optional[Callable] = None,
      regularizer: Optional[Callable] = None,
  ):
      """
      Train for `epoch` epochs on a Dataset `dataset.
  
      Args:
          dataset (Dataset): Pytorch Dataset to be trained on.
          optimizer (optim.Optimizer): Optimizer to use.
          batch_size (int): The batch size used in the DataLoader.
          epoch (int): Number of epoch to train for.
          use_cuda (bool): Use cuda or not.
          workers (int): Number of workers for the multiprocessing.
          collate_fn (Optional[Callable]): The collate function to use.
          regularizer (Optional[Callable]): The loss regularization for training.
  
      Returns:
          The training history.
      """
      dataset_size = len(dataset)
      self.train()
      self.set_dataset_size(dataset_size)
      history = []
      log.info("Starting training", epoch=epoch, dataset=dataset_size)
      collate_fn = collate_fn or default_collate
      for _ in range(epoch):
          self._reset_metrics("train")
          for data, target, *_ in DataLoader(
              dataset,
              batch_size,
              True,
              num_workers=workers,
              collate_fn=collate_fn,
          ):
              _ = self.train_on_batch(
                  data, target, optimizer, use_cuda, regularizer
              )
          history.append(self.get_metrics("train")["train_loss"])
  
      optimizer.zero_grad()  # Assert that the gradient is flushed.
      log.info(
          "Training complete",
          train_loss=self.get_metrics("train")["train_loss"],
      )
      self.active_step(dataset_size, self.get_metrics("train"))
      return history
  
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 2 for calling the ModelWrapper obj's underlying member method code with example usage: obj.train_on_dataset(dataset, optimizer, batch_size, epoch, use_cuda, workers, collate_fn, regularizer) and returns: history