class_docstr: Pytorch Lightning class which adds methods to perform
  active learning.
docstr: Predict on batch using MC inference `I` times.
  `I` is defined in the hparams property.
  Args:
      data (Tensor): Data to feed to the model.
      batch_idx (int): Batch index.
      dataloader_idx: Index of the current dataloader (not used)
  
  Returns:
      Models predictions stacked `I` times on the last axis.
  
  Notes:
      If `hparams.replicate_in_memeory` is True, we will stack inputs I times.
      This might create OoM errors. In that case, set it to False.
docstr_tok: ["Predict","on","batch","using","MC","inference","`","I","`","times",".","`","I","`","is","defined","in","the","hparams","property",".","Args",":","data","(","Tensor",")",":","Data","to","feed","to","the","model",".","batch_idx","(","int",")",":","Batch","index",".","dataloader_idx",":","Index","of","the","current","dataloader","(","not","used",")","Returns",":","Models","predictions","stacked","`","I","`","times","on","the","last","axis",".","Notes",":","If","`","hparams.replicate_in_memeory","`","is","True",",","we","will","stack","inputs","I","times",".","This","might","create","OoM","errors",".","In","that","case",",","set","it","to","False","."]
code: def predict_step(
      self, batch, batch_idx, dataloader_idx: Optional[int] = None
  ):
      """Predict on batch using MC inference `I` times.
      `I` is defined in the hparams property.
      Args:
          data (Tensor): Data to feed to the model.
          batch_idx (int): Batch index.
          dataloader_idx: Index of the current dataloader (not used)
  
      Returns:
          Models predictions stacked `I` times on the last axis.
  
      Notes:
          If `hparams.replicate_in_memeory` is True, we will stack inputs I times.
          This might create OoM errors. In that case, set it to False.
      """
      # Get the input only.
      x, _ = batch
      # Perform Monte-Carlo Inference fro I iterations.
      out = mc_inference(
          self, x, self.hparams.iterations, self.hparams.replicate_in_memory  # type: ignore
      )  # type: ignore
      return out
  
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 ActiveLightningModule obj's underlying member method code with example usage: obj.predict_step(batch, batch_idx, dataloader_idx) and returns: out