def METHOD_NAME(self, trainer, pl_module):
    if self.data_sparsifier_state_dict is None:
        return  # probably first epoch
    # load the existing config for each data
    self.data_sparsifier.load_state_dict(self.data_sparsifier_state_dict)