code: def test_lr_callback_batch_steps_correctly_fallback(
      self,
      classifier_module,
      classifier_data,
      policy,
      kwargs,
  ):
      batch_size = 100
      max_epochs = 2
  
      X, y = classifier_data
      num_examples = len(X)
  
      lr_policy = LRScheduler(policy, **kwargs)
      net = NeuralNetClassifier(
          classifier_module(),
          max_epochs=max_epochs,
          batch_size=batch_size,
          callbacks=[lr_policy],
      )
      net.fit(X, y)
  
      # Removes batch count information in the last two epochs
      for i in range(max_epochs):
          del net.history[i]["train_batch_count"]
          del net.history[i]["valid_batch_count"]
      net.partial_fit(X, y)
  
      total_iterations_per_epoch = num_examples / batch_size
  
      # batch_counts were removed thus the total iterations of the last
      # epoch is used
      total_iterations_fit_run = total_iterations_per_epoch * max_epochs
  
      # 80% of sample used for training by default
      total_iterations_partial_fit_run = (
          0.8 * total_iterations_per_epoch * max_epochs
      )
  
      # called fit AND partial_fit
      total_iterations = (
          total_iterations_fit_run + total_iterations_partial_fit_run
      )
      # Failback to using both valid and training batches counts on
      # second run
      expected = int(total_iterations)
      # pylint: disable=protected-access
      assert lr_policy.batch_idx_ == expected
  
imports: from unittest.mock import Mock
  import numpy
  import pytest
  from sklearn.base import clone
  from torch.optim import SGD
  from torch.optim.lr_scheduler import CosineAnnealingLR
  from torch.optim.lr_scheduler import ExponentialLR
  from torch.optim.lr_scheduler import LambdaLR
  from torch.optim.lr_scheduler import MultiStepLR
  from torch.optim.lr_scheduler import ReduceLROnPlateau
  from torch.optim.lr_scheduler import StepLR
  from torch.optim.lr_scheduler import CyclicLR
  from skorch import NeuralNetClassifier
  from skorch.callbacks.lr_scheduler import WarmRestartLR, LRScheduler
label_desc: Use image node_id 6 for calling the TestLRCallbacks obj's underlying member method code with example usage: obj.test_lr_callback_batch_steps_correctly_fallback(classifier_module, classifier_data, policy, kwargs) without return types