def METHOD_NAME(self, step: int, epoch: int) -> None:
    """
    Called after the training loop completes.
    This is the last method that is called, so any cleanup can be done in this method.
    """
    pass