code: def test_active_learning_training(
      simple_datamodule, initial_num_labels, query_size
  ):
      seed_everything(42)
  
      if initial_num_labels == 0:
          with pytest.warns(
              UserWarning,
              match="No labels provided for the initial step",
          ):
              active_learning_dm = ActiveLearningDataModule(
                  simple_datamodule,
                  initial_num_labels=initial_num_labels,
                  query_size=query_size,
                  val_split=0.5,
              )
      else:
          active_learning_dm = ActiveLearningDataModule(
              simple_datamodule,
              initial_num_labels=initial_num_labels,
              query_size=query_size,
              val_split=0.5,
          )
  
      head = nn.Sequential(
          nn.Dropout(p=0.1),
          nn.Linear(512, active_learning_dm.num_classes),
      )
  
      model = ImageClassifier(
          backbone="resnet18",
          head=head,
          num_classes=active_learning_dm.num_classes,
      )
      trainer = flash.Trainer(max_epochs=3, num_sanity_val_steps=0)
      active_learning_loop = ActiveLearningLoop(
          label_epoch_frequency=1, inference_iteration=3
      )
      active_learning_loop.connect(trainer.fit_loop)
      trainer.fit_loop = active_learning_loop
  
      trainer.finetune(
          model, datamodule=active_learning_dm, strategy="no_freeze"
      )
      # Check that all metrics are logged
      assert all(
          any(
              m in log_met
              for log_met in active_learning_loop.trainer.logged_metrics
          )
          for m in ("train", "val", "test")
      )
  
      # Check that the weights has changed for both module.
      classifier = (
          active_learning_loop._lightning_module.adapter.parameters()
      )
      mc_inference = (
          active_learning_loop.inference_model.parent_module.parameters()
      )
      assert all(
          torch.equal(p1, p2)
          for p1, p2 in zip(classifier, mc_inference)
      )
  
      if initial_num_labels == 0:
          assert len(active_learning_dm._dataset) == 15
      else:
          assert len(active_learning_dm._dataset) == 20
      assert active_learning_loop.progress.total.completed == 3
      labelled = active_learning_loop.state_dict()["state_dict"][
          "datamodule_state_dict"
      ]["labelled"]
      if _BAAL_GREATER_EQUAL_1_5_2:
          labelled = labelled > 0
      assert isinstance(labelled, np.ndarray)
  
      # Check that we iterate over the actual pool and that shuffle is disabled.
      assert len(active_learning_dm.predict_dataloader()) == math.ceil(
          (~labelled).sum() / simple_datamodule.batch_size
      )
      assert isinstance(
          active_learning_dm.predict_dataloader().sampler,
          SequentialSampler,
      )
  
      if initial_num_labels == 0:
          assert len(active_learning_dm.val_dataloader()) == 4
      else:
          # in the second scenario we have more labelled data!
          assert len(active_learning_dm.val_dataloader()) == 5
  
file_path: lightning-flash/tests/image/classification/test_active_learning.py
imports: import math
  from pathlib import Path
  import flash
  import numpy
  import pytest
  import torch
  from flash.core.utilities.imports import _BAAL_AVAILABLE, _BAAL_GREATER_EQUAL_1_5_2, _TOPIC_IMAGE_AVAILABLE
  from flash.image import ImageClassificationData, ImageClassifier
  from flash.image.classification.integrations.baal import ActiveLearningDataModule, ActiveLearningLoop
  from pytorch_lightning import seed_everything
  from torch import nn
  from torch.utils.data import SequentialSampler
  from tests.image.classification.test_data import _rand_image
label_desc: Use image node_id 2 for calling a global function with example usage: test_active_learning_training(simple_datamodule, initial_num_labels, query_size) without return types