Early stopping + checkpoint key
Consider updating how we condition early stopping or checkpoint
return {'early_stop_on': mse_loss, 'checkpoint_on': other_metric}
Instead of:
# only if val_loss is present
return {'val_loss': val_loss}