# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/callback.tracker.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['TerminateOnTrainNaN', 'SaveModelAtEnd']

# Cell
#nbdev_comment from __future__ import annotations

from fastai.callback.core import Callback, CancelFitException
from fastai.callback.tracker import SaveModelCallback, TrackerCallback

from ..imports import *

# Cell
class TerminateOnTrainNaN(Callback):
    "A `Callback` that terminates training if the training loss is NaN and ignores valid loss."
    order,run_valid=-9,False
    def after_batch(self):
        "Test if `last_loss` is NaN and interrupts training."
        if torch.isinf(self.loss) or torch.isnan(self.loss): raise CancelFitException

# Cell
class SaveModelAtEnd(SaveModelCallback):
    "A `SaveModelCallback` which only saves the model at the end so loggers can find it."
    order = TrackerCallback.order+1
    def __init__(self,
        fname='model', # Model filename
        with_opt=False # Include optimizer state
    ):
        # keep track of file path for loggers
        self.last_saved_path = None
        store_attr()

    def before_fit(self):
        pass

    def after_epoch(self):
        pass

    def after_fit(self, **kwargs):
        self.last_saved_path = self.learn.save(f'{self.fname}', with_opt=self.with_opt)