code: def __init__(
      self,
      active_dataset: ActiveLearningDataset,
      batch_size=1,
      **kwargs,
  ):
      super().__init__(**kwargs)
      self.active_dataset = active_dataset
      self.batch_size = batch_size
  
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 1 to create a new BaaLDataModule object from inherited base classes: LightningDataModule with example: obj = BaaLDataModule(active_dataset, batch_size)