class_docstr: Object that perform the training and active learning iteration.
  
  Args:
      dataset (ActiveLearningDataset): Dataset with some sample already labelled.
      heuristic (Heuristic): Heuristic from baal.active.heuristics.
      query_size (int): Number of sample to label per step.
      max_sample (int): Limit the number of sample used (-1 is no limit).
      **kwargs: Parameters forwarded to `get_probabilities`
          and to pytorch_ligthning Trainer.__init__
docstr: For documentation, see `predict_on_dataset_generator`
docstr_tok: ["For","documentation",",","see","`","predict_on_dataset_generator","`"]
code: def predict_on_dataset(
      self, model=None, dataloader=None, *args, **kwargs
  ):
      "For documentation, see `predict_on_dataset_generator`"
      preds = list(self.predict_on_dataset_generator(model, dataloader))
  
      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
  import types
  import warnings
  from collections.abc import Sequence
  from typing import Dict, Any, Optional
  import numpy
  import structlog
  import torch
  from pytorch_lightning import Trainer, Callback, LightningDataModule, LightningModule
  from pytorch_lightning.accelerators import CUDAAccelerator
  from torch.utils.data import DataLoader
  from tqdm import tqdm
  from baal.active import ActiveLearningDataset
  from baal.active.heuristics import heuristics
  from baal.modelwrapper import mc_inference
  from baal.utils.cuda_utils import to_cuda
  from baal.utils.iterutils import map_on_tensor
label_desc: Use image node_id 2 for calling the BaalTrainer obj's underlying member method code with example usage: obj.predict_on_dataset(model, dataloader) and returns: unknown, np