def __init__(self, checkpoint_io: Optional[CheckpointIO] = None) -> None:
self._model: Optional[Module] = None
self._results: Optional[Union[_EVALUATE_OUTPUT,
_PREDICT_OUTPUT]] = None
checkpoint_io = checkpoint_io if checkpoint_io is not None else TorchCheckpointIO(
)
self._checkpoint_io = checkpoint_io
def checkpoint_io(self) -> CheckpointIO:
if self._checkpoint_io is None:
self._checkpoint_io = TorchCheckpointIO()
return self._checkpoint_io
def checkpoint_io(self) -> CheckpointIO:
return self._checkpoint_io if self._checkpoint_io is not None else TorchCheckpointIO()