code: def after_epoch(self, epoch_data: EpochData, **kwargs) -> None:
          
  
          if self._after_epoch:
              self._check_nan(epoch_data)
docs: If initialized to check after each epoch, stop the training once the epoch data contains a monitored
          variable equal to NaN.
  
          :param epoch_data: epoch data to be checked