Using a waiter to wait instead of attach

aws · Jun 14, 2022 · ffce787 · ffce787
1 parent 4873fb0
commit ffce787
Showing 1 changed file with 13 additions and 0 deletions.
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
@@ -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)"