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: Confirms that the mask is squashed after the training ends
  This is achieved by making sure that each parameter in the internal container
  are not parametrized.
docstr_tok: ["Confirms","that","the","mask","is","squashed","after","the","training","ends","This","is","achieved","by","making","sure","that","each","parameter","in","the","internal","container","are","not","parametrized","."]
code: def _check_on_train_end(self, pl_module, callback):
      """Confirms that the mask is squashed after the training ends
      This is achieved by making sure that each parameter in the internal container
      are not parametrized.
      """
      callback.on_train_end(42, pl_module)
  
      # check that the masks have been squashed
      for name, _ in pl_module.model.named_parameters():
          valid_name = _get_valid_name(name)
          assert not is_parametrized(
              callback.data_sparsifier._continer, valid_name
          )
  
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 5 for calling the TestTrainingAwareCallback obj's underlying member method code with example usage: obj._check_on_train_end(pl_module, callback) without return types