code: def test_weight_assignment(self):
      params = list(self.wrapper.model.parameters())
      self.wrapper.train_on_dataset(
          self.dataset, self.optim, 32, 1, False
      )
      assert all(
          [
              k is v
              for k, v in zip(
                  params, self.optim.param_groups[0]["params"]
              )
          ]
      )
  
      self.calibrator.calibrate(
          self.dataset, self.dataset, 32, 1, False, True
      )
      assert all(
          [
              k is v
              for k, v in zip(
                  self.wrapper.model.parameters(),
                  self.optim.param_groups[0]["params"],
              )
          ]
      )
  
      # Check that we can train the original model
      before_params = list(
          map(lambda x: x.clone(), self.wrapper.model.parameters())
      )
      self.wrapper.train_on_dataset(
          self.dataset, self.optim, 10, 2, False
      )
      after_params = list(
          map(lambda x: x.clone(), self.wrapper.model.parameters())
      )
      assert not all(
          [
              np.allclose(i.detach(), j.detach())
              for i, j in zip(before_params, after_params)
          ]
      )
  
      # Check that the parameters are still tied.
      calib_params = list(
          map(
              lambda x: x.clone(),
              _get_first_module(
                  self.calibrator.wrapper.model
              ).parameters(),
          )
      )
      assert all(
          [
              np.allclose(i.detach(), j.detach())
              for i, j in zip(calib_params, after_params)
          ]
      )
  
imports: import unittest
  from unittest.mock import Mock
  import numpy
  import pytest
  import torch
  from torch import nn
  from torch.utils.data import Dataset
  from baal.calibration import DirichletCalibrator
  from baal.modelwrapper import ModelWrapper
label_desc: Use image node_id 5 for calling the CalibrationTest obj's underlying member method code with example usage: obj.test_weight_assignment() without return types