code: def test_train_on_batch(self):
      self.wrapper.train()
      old_param = list(
          map(lambda x: x.clone(), self.model.parameters())
      )
      input, target = torch.randn([1, 3, 10, 10]), torch.randn(1, 1)
      self.wrapper.train_on_batch(input, target, self.optim)
      new_param = list(
          map(lambda x: x.clone(), self.model.parameters())
      )
  
      assert any(
          [
              not torch.allclose(i, j)
              for i, j in zip(old_param, new_param)
          ]
      )
  
      # test reset weights properties
      linear_weights = list(self.wrapper.model.named_children())[3][
          1
      ].weight.clone()
      conv_weights = list(self.wrapper.model.named_children())[0][
          1
      ].weight.clone()
      self.wrapper.reset_fcs()
      linear_new_weights = list(self.wrapper.model.named_children())[3][
          1
      ].weight.clone()
      conv_new_weights = list(self.wrapper.model.named_children())[0][
          1
      ].weight.clone()
      assert all(
          [
              not torch.allclose(i, j)
              for i, j in zip(linear_new_weights, linear_weights)
          ]
      )
      assert all(
          [
              torch.allclose(i, j)
              for i, j in zip(conv_new_weights, conv_weights)
          ]
      )
  
      self.wrapper.reset_all()
      conv_next_new_weights = list(self.wrapper.model.named_children())[
          0
      ][1].weight.clone()
      assert all(
          [
              not torch.allclose(i, j)
              for i, j in zip(conv_new_weights, conv_next_new_weights)
          ]
      )
  
imports: import math
  import unittest
  from unittest.mock import Mock
  import numpy
  import pytest
  import torch
  import torchmetrics
  from torch import nn
  from torch.utils.data import Dataset, DataLoader
  from baal.modelwrapper import ModelWrapper, mc_inference
  from baal.utils.metrics import ClassificationReport
label_desc: Use image node_id 2 for calling the ModelWrapperTest obj's underlying member method code with example usage: obj.test_train_on_batch() without return types