Store the notebook and stop run

def on_train_end(self, **kwargs: Any) -> None:  
        "Store the notebook and stop run"
        self.client.log_artifact(run_id=self.run, local_path=self.nb_path)
        self.client.set_terminated(run_id=self.run)