class_docstr: Wrapper that support ensembles with the same architecture.
  
  Args:
      model (nn.Module): A Model.
      criterion (Callable): Loss function
  
  Notes:
      If you're looking to use ensembles for non-deep models, see our sklearn tutorial:
       baal.readthedocs.io/en/latest/notebooks/sklearn_tutorial.html
docstr: Get the model's prediction on a batch.
  
  Args:
      data (Tensor): The model input.
      iterations (int): NOT USED
      cuda (bool): Use CUDA or not.
  
  Returns:
      Tensor, the loss computed from the criterion.
              shape = {batch_size, nclass, n_iteration}.
  
  Raises:
      Raises ValueError when no checkpoint are logged.
docstr_tok: ["Get","the","model","'s","prediction","on","a","batch",".","Args",":","data","(","Tensor",")",":","The","model","input",".","iterations","(","int",")",":","NOT","USED","cuda","(","bool",")",":","Use","CUDA","or","not",".","Returns",":","Tensor",",","the","loss","computed","from","the","criterion",".","shape","=","{","batch_size",",","nclass",",","n_iteration","}",".","Raises",":","Raises","ValueError","when","no","checkpoint","are","logged","."]
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): NOT USED
          cuda (bool): Use CUDA or not.
  
      Returns:
          Tensor, the loss computed from the criterion.
                  shape = {batch_size, nclass, n_iteration}.
  
      Raises:
          Raises ValueError when no checkpoint are logged.
      """
      if len(self._weights) == 0:
          raise ValueError("No checkpoint!")
  
      return ensemble_prediction(
          model=self.model, data=data, weights=self._weights, cuda=cuda
      )
  
file_path: baal/baal/ensemble.py
imports: from copy import deepcopy
  from typing import Dict, List, OrderedDict
  import torch
  from torch import nn, Tensor
  from baal import ModelWrapper
  from baal.modelwrapper import _stack_preds
  from baal.utils.cuda_utils import to_cuda
label_desc: Use image node_id 4 for calling the EnsembleModelWrapper obj's underlying member method code with example usage: obj.predict_on_batch(data, iterations, cuda) and returns: ensemble_prediction