From ffce78787904c1730cbfeec1a02ac47173a6c807 Mon Sep 17 00:00:00 2001 From: Loki Date: Tue, 14 Jun 2022 17:29:53 +0000 Subject: [PATCH] Using a waiter to wait instead of attach --- .../single_gpu_single_node/vision-transformer.ipynb | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/sagemaker-training-compiler/tensorflow/single_gpu_single_node/vision-transformer.ipynb b/sagemaker-training-compiler/tensorflow/single_gpu_single_node/vision-transformer.ipynb index d5206afd57..6de8d245fe 100644 --- a/sagemaker-training-compiler/tensorflow/single_gpu_single_node/vision-transformer.ipynb +++ b/sagemaker-training-compiler/tensorflow/single_gpu_single_node/vision-transformer.ipynb @@ -331,6 +331,19 @@ "scrolled": true }, "outputs": [], + "source": [ + "waiter = sess.sagemaker_client.get_waiter(\"training_job_completed_or_stopped\")\n", + "\n", + "waiter.wait(TrainingJobName=native_estimator.latest_training_job.name)\n", + "waiter.wait(TrainingJobName=optimized_estimator.latest_training_job.name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f9217c89", + "metadata": {}, + "outputs": [], "source": [ "native_estimator = TensorFlow.attach(native_estimator.latest_training_job.name)\n", "optimized_estimator = TensorFlow.attach(optimized_estimator.latest_training_job.name)"