def _getCheckpointCallback(self, version: int = None, stage=None): model_dir = self.paths.get("model_dir", f'model/{DATE}') model_dir = os.path.join(model_dir, self.name) if version or version == 0: version = f"version_{version}" if isinstance(version, int) else version model_dir = os.path.join(model_dir, version) exist = os.path.exists(model_dir) and os.listdir(model_dir) if stage == 'test' and not exist: raise FileNotFoundError(model_dir) elif stage is None and exist: raise FileExistsError(model_dir) self.model_dir = model_dir checkpoint_callback = ModelCheckpoint( monitor="err/B-M/validation", dirpath=model_dir, filename="best", save_last=True, mode="min", ) checkpoint_callback.FILE_EXTENSION = ".pth" checkpoint_callback.best_model_path = model_dir checkpoint_callback.CHECKPOINT_NAME_LAST = "latest" return checkpoint_callback