def METHOD_NAME(self, model, checkpoint_path):
    if self._load_checkpoint_fn is not None:
        self._load_checkpoint_fn(model, checkpoint_path)
    else:
        model.load_weights(checkpoint_path)