code: def test_reset_callback_resets_weights(a_data_module):
      def reset_fcs(model):
          """Reset all torch.nn.Linear layers."""
  
          def reset(m):
              if isinstance(m, torch.nn.Linear):
                  m.reset_parameters()
  
          model.apply(reset)
  
      model = vgg16()
      trainer = BaalTrainer(
          dataset=a_data_module.active_dataset,
          max_epochs=3,
          default_root_dir="/tmp",
      )
      trainer.fit_loop.current_epoch = 10
      initial_weights = copy.deepcopy(model.state_dict())
      initial_params = copy.deepcopy(list(model.parameters()))
      callback = ResetCallback(initial_weights)
      # Modify the params
      reset_fcs(model)
      new_params = model.parameters()
      assert not all(
          torch.eq(p1, p2).all()
          for p1, p2 in zip(initial_params, new_params)
      )
      callback.on_train_start(trainer, model)
      new_params = model.parameters()
      assert all(
          torch.eq(p1, p2).all()
          for p1, p2 in zip(initial_params, new_params)
      )
      assert trainer.current_epoch == 0
  
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 8 for calling a global function with example usage: test_reset_callback_resets_weights(a_data_module) without return types