code: def test_step(self):
      # For n=5, dt=2, there will be totally 10 steps between s_0 and s_f, starting from t_0
      model = self._make_model()
      sparsifier, scheduler = self._make_scheduler(
          model=model,
          initially_zero=True,
          init_t=3,
          delta_t=2,
          total_t=5,
      )
  
      scheduler.step()
      scheduler.step()
      self.assertEqual(
          scheduler._step_count,
          3,
          msg="Scheduler step_count is expected to increment",
      )
      # Value before t_0 is supposed to be 0
      self.assertEqual(
          self._get_sparsity_levels(sparsifier),
          scheduler._make_sure_a_list(0.0),
          msg="Scheduler step updating the sparsity level before t_0",
      )
  
      scheduler.step()  # Step = 3  =>  sparsity = initial_sparsity
      self.assertEqual(
          self._get_sparsity_levels(sparsifier),
          scheduler._make_sure_a_list(self.initial_sparsity),
          msg="Sparsifier is not reset to initial sparsity at the first step",
      )
  
      scheduler.step()  # Step = 4  =>  sparsity ~ [0.3, 0.2]
      self.assertEqual(
          self._get_sparsity_levels(sparsifier, 1),
          [0.3, 0.2],
          msg="Sparsity level is not set correctly after the first step",
      )
  
      current_step = scheduler._step_count - scheduler.init_t[0] - 1
      more_steps_needed = (
          scheduler.delta_t[0] * scheduler.total_t[0] - current_step
      )
      for _ in range(
          more_steps_needed
      ):  # More steps needed to final sparsity level
          scheduler.step()
      self.assertEqual(
          self._get_sparsity_levels(sparsifier),
          self.sorted_sparse_levels,
          msg="Sparsity level is not reaching the target level afer delta_t * n steps ",
      )
  
imports: from torch import nn
  from torch.ao.pruning import WeightNormSparsifier
  from torch.ao.pruning import BaseScheduler, LambdaSL, CubicSL
  from torch.testing._internal.common_utils import TestCase
  import warnings
label_desc: Use image node_id 6 for calling the TestCubicScheduler obj's underlying member method code with example usage: obj.test_step() without return types