From 6f86c68cfde3a94db3f79a9627222331fac3d53f Mon Sep 17 00:00:00 2001
From: Krishna Chaitanya Koppolu <71738025+kkoppolu1@users.noreply.github.com>
Date: Fri, 9 Oct 2020 17:30:25 -0400
Subject: [PATCH] Verify sagemaker SDK version (#1606)

---
 .../pytorch_torchvision_neo.ipynb             | 19 +++++++++++++++++++
 .../pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb   | 19 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb b/sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb
index 691647bf9d..f340504edd 100644
--- a/sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb
+++ b/sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb
@@ -23,6 +23,13 @@
     "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.4.0 torchvision==0.5.0"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### SageMaker SDK >= 2.0 is required for this notebook"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -32,6 +39,18 @@
     "!~/anaconda3/envs/pytorch_p36/bin/pip install --upgrade sagemaker"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sagemaker\n",
+    "current_version = sagemaker.__version__\n",
+    "if current_version.split('.')[0] == '1':\n",
+    "    raise Exception(\"Please upgrade SageMaker SDK by running the above code cell and restart the kernel\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb b/sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb
index 0631520344..60fbc5ffb4 100644
--- a/sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb
+++ b/sagemaker_neo_compilation_jobs/pytorch_vgg19_bn/pytorch-vgg19-bn.ipynb
@@ -27,6 +27,13 @@
     "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.4.0 torchvision==0.5.0"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### SageMaker SDK >= 2.0 is required for this notebook"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -36,6 +43,18 @@
     "!~/anaconda3/envs/pytorch_p36/bin/pip install --upgrade sagemaker"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sagemaker\n",
+    "current_version = sagemaker.__version__\n",
+    "if current_version.split('.')[0] == '1':\n",
+    "    raise Exception(\"Please upgrade SageMaker SDK by running the above code cell and restart the kernel\")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},