code: def test_train_on_batch(self):
      self.wrapper.train()
      old_param = list(
          map(lambda x: x.clone(), self.model.parameters())
      )
      input, target = [
          torch.stack(v)
          for v in zip(*(self.dataset[0], self.dataset[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)
          ]
      )
  
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 ModelWrapperMultiOutMultiInTest obj's underlying member method code with example usage: obj.test_train_on_batch() without return types