callback that terminates training when a nan loss is encountered def onbatchendself batch logsnone logs logs or loss logs getloss if loss is not none if np isnanloss or np isinfloss ioutils printmsg fbatch batch invalid loss terminating training self model stoptraining true callback that terminates training when a nan loss is encountered
import numpy as np

from keras.api_export import keras_export
from keras.callbacks.callback import Callback
from keras.utils import io_utils


@keras_export("keras.callbacks.TerminateOnNaN")
class TerminateOnNaN(Callback):
    

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        loss = logs.get("loss")
        if loss is not None:
            if np.isnan(loss) or np.isinf(loss):
                io_utils.print_msg(
                    f"Batch {batch}: Invalid loss, terminating training"
                )
                self.model.stop_training = True