code: def test_on_load_checkpoint(a_data_module, a_dataset, hparams):
      ckpt = {}
      _ = a_data_module.on_save_checkpoint(ckpt)
      assert "active_dataset" in ckpt
      active_set_2 = ActiveLearningDataset(a_dataset)
      data_mdoule_2 = MyDataModule(active_set_2, hparams["batch_size"])
      on_load_chkp = data_mdoule_2.on_load_checkpoint(ckpt)
      assert len(a_data_module.active_dataset) == len(active_set_2)
  
imports: import copy
  from dataclasses import dataclass, asdict
  import numpy
  import pytest
  import torch
  from baal.active import ActiveLearningDataset
  from baal.utils.pytorch_lightning import ActiveLightningModule, BaalTrainer, ResetCallback, BaaLDataModule
  from pytorch_lightning import LightningModule
  from torch.utils.data import Dataset, DataLoader
  from torchvision.models import vgg16
label_desc: Use image node_id 6 for calling a global function with example usage: test_on_load_checkpoint(a_data_module, a_dataset, hparams) without return types