class_docstr: Class to test in-training version of lightning callback
  Simulates model training and makes sure that each hook is doing what is expected
class_docstr_tok: ["Class","to","test","in-training","version","of","lightning","callback","Simulates","model","training","and","makes","sure","that","each","hook","is","doing","what","is","expected"]
docstr: Checks the following -
  1. sparsity is correctly being achieved after .step()
  2. scheduler and data_sparsifier sparsity levels are in-line
docstr_tok: ["Checks","the","following","-","1.","sparsity","is","correctly","being","achieved","after",".step","(",")","2.","scheduler","and","data_sparsifier","sparsity","levels","are","in-line"]
code: def _check_on_train_epoch_end(self, pl_module, callback):
      """Checks the following -
      1. sparsity is correctly being achieved after .step()
      2. scheduler and data_sparsifier sparsity levels are in-line
      """
      callback.on_train_epoch_end(42, pl_module)
      data_scheduler = callback.data_scheduler
      base_sl = data_scheduler.base_param
  
      for name, _ in pl_module.model.named_parameters():
          valid_name = _get_valid_name(name)
          mask = callback.data_sparsifier.get_mask(name=valid_name)
  
          # check sparsity levels
          assert (
              1.0 - mask.float().mean()
          ) > 0  # some sparsity level achieved
  
          last_sl = data_scheduler.get_last_param()
          last_epoch = data_scheduler.last_epoch
  
          # check sparsity levels of scheduler
          log_last_sl = math.log(last_sl[valid_name])
          log_actual_sl = math.log(
              base_sl[valid_name] * (data_scheduler.gamma**last_epoch)
          )
          assert log_last_sl == log_actual_sl
  
imports: from torch.ao.pruning._experimental.data_sparsifier.data_norm_sparsifier import DataNormSparsifier
  from torch.ao.pruning._experimental.data_scheduler.base_data_scheduler import BaseDataScheduler
  import torch
  import torch.nn
  from typing import List
  from torch.ao.pruning._experimental.data_sparsifier.lightning.callbacks.data_sparsity import PostTrainingDataSparsity, TrainingAwareDataSparsity
  from torch.ao.pruning._experimental.data_sparsifier.lightning.callbacks._data_sparstity_utils import _get_valid_name
  from torch.ao.pruning._experimental.data_sparsifier.base_data_sparsifier import SUPPORTED_TYPES
  from torch.testing._internal.common_utils import TestCase
  from torch.testing._internal.common_utils import run_tests
  import importlib
  import unittest
  import warnings
  import math
  from torch.nn.utils.parametrize import is_parametrized
label_desc: Use image node_id 4 for calling the TestTrainingAwareCallback obj's underlying member method code with example usage: obj._check_on_train_epoch_end(pl_module, callback) without return types