class InterruptAtEpoch(keras.callbacks.Callback):
    def __init__(self, epoch_to_stop):
        super().__init__()
        self.epoch_to_stop = epoch_to_stop

    def on_epoch_begin(self, epoch, logs=None):
        if epoch == self.epoch_to_stop:
            raise RuntimeError(f"simulated interruption before epoch {epoch + 1}")