    def on_train_end(self, logs: Dict) -> None:
        """
        Hook function invoked after train stage.

        Args:
            logs (`Dict`):
               Kv store to save and load info.
        """
        for cb in self.callbacks:
            cb.on_train_end(logs)
Hook function invoked after train stage.

Args:
logs (`Dict`):
Kv store to save and load info.
towhee/trainer/callback.py