code: def test_calibration_integration():
      transform_pipeline = Compose([Resize((64, 64)), ToTensor()])
      cifar10_train = DummyDataset(transform_pipeline)
      cifar10_test = DummyDataset(transform_pipeline)
  
      # we don't create different trainset for calibration since the goal is not
      # to calibrate
      al_dataset = ActiveLearningDataset(
          cifar10_train,
          pool_specifics={"transform": transform_pipeline},
      )
      al_dataset.label_randomly(10)
      use_cuda = False
      model = vgg.vgg16(pretrained=False, num_classes=10)
  
      criterion = nn.CrossEntropyLoss()
      optimizer = optim.SGD(
          model.parameters(),
          lr=0.001,
          momentum=0.9,
          weight_decay=0.0005,
      )
  
      wrapper = ModelWrapper(model, criterion)
      calibrator = DirichletCalibrator(
          wrapper=wrapper, num_classes=10, lr=0.001, reg_factor=0.01
      )
  
      for step in range(2):
          wrapper.train_on_dataset(
              al_dataset,
              optimizer=optimizer,
              batch_size=10,
              epoch=1,
              use_cuda=use_cuda,
              workers=0,
          )
  
          wrapper.test_on_dataset(
              cifar10_test, batch_size=10, use_cuda=use_cuda, workers=0
          )
  
          before_calib_param = list(
              map(lambda x: x.clone(), wrapper.model.parameters())
          )
  
          calibrator.calibrate(
              al_dataset,
              cifar10_test,
              batch_size=10,
              epoch=5,
              use_cuda=use_cuda,
              double_fit=False,
              workers=0,
          )
  
          after_calib_param = list(
              map(lambda x: x.clone(), model.parameters())
          )
  
          assert all(
              [
                  np.allclose(i.detach(), j.detach())
                  for i, j in zip(before_calib_param, after_calib_param)
              ]
          )
  
          assert len(list(wrapper.model.modules())) < len(
              list(calibrator.calibrated_model.modules())
          )
  
imports: import os
  import numpy
  import pytest
  from PIL import Image
  from torch import nn
  from torch import optim
  from torch.utils.data import Dataset
  from torchvision.models import vgg
  from torchvision.transforms import ToTensor, Resize, Compose
  from baal.active import ActiveLearningDataset
  from baal.active import ActiveLearningLoop
  from baal.active import heuristics
  from baal.modelwrapper import ModelWrapper
  from baal.calibration import DirichletCalibrator
label_desc: Use image node_id 2 for calling a global function with example usage: test_calibration_integration() without return types