fix: use image_uris module for retrieval

aws · Nov 6, 2020 · b73e39d · b73e39d
1 parent 6553d0d
commit b73e39d
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/sagemaker-python-sdk/scikit_learn_randomforest/Sklearn_on_SageMaker_end2end.ipynb b/sagemaker-python-sdk/scikit_learn_randomforest/Sklearn_on_SageMaker_end2end.ipynb
@@ -331,11 +331,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sagemaker.fw_registry import default_framework_uri\n",
+    "from sagemaker import image_uris\n",
     "\n",
-    "training_image = default_framework_uri(\n",
-    "    'scikit-learn', region, '{}-cpu-py3'.format(FRAMEWORK_VERSION))\n",
-    "print(training_image)"
+    "\n",
+    "training_image = image_uris.retrieve(\n",
+    "    framework=\"sklearn\",\n",
+    "    region=region,\n",
+    "    version=FRAMEWORK_VERSION,\n",
+    "    py_version=\"py3\",\n",
+    "    instance_type=\"ml.c5.xlarge\"\n",
+    ")\n",
+    "print(training_image)\n"
    ]
   },
   {
@@ -428,7 +434,7 @@
     "    metric_definitions=[\n",
     "        {'Name': 'median-AE',\n",
     "         'Regex': \"AE-at-50th-percentile: ([0-9.]+).*$\"}],  # extract tracked metric from logs with regexp \n",
-    "    max_jobs=20,\n",
+    "    max_jobs=10,\n",
     "    max_parallel_jobs=2)"
    ]
   },