Upgrade Pytorch version to 1.2.0 in Neo notebook (#1098)

* set pytorch version for neo sagemaker notebook

* Upgrade pytorch version to 1.2.0 for neo notebook

Co-authored-by: Hui Wang <[email protected]>

Author: melty-chocolate

sagemaker_neo_compilation_jobs/pytorch_torchvision/pytorch_torchvision_neo.ipynb

@@ -17,10 +17,12 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "metadata": {},
+   "metadata": {
+    "scrolled": true
+   },
    "outputs": [],
    "source": [
-    "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.0"
+    "!~/anaconda3/envs/pytorch_p36/bin/pip install torch==1.2.0"
    ]
   },
   {
@@ -48,7 +50,9 @@
     "import tarfile\n",
     "\n",
     "resnet18 = models.resnet18(pretrained=True)\n",
-    "torch.save(resnet18, 'model.pth')\n",
+    "input_shape = [1,3,224,224]\n",
+    "trace = torch.jit.trace(resnet18.float().eval(), torch.zeros(input_shape).float())\n",
+    "trace.save('model.pth')\n",
     "\n",
     "with tarfile.open('model.tar.gz', 'w:gz') as f:\n",
     "    f.add('model.pth')"
@@ -94,7 +98,7 @@
     "data_shape = '{\"input0\":[1,3,224,224]}'\n",
     "target_device = 'ml_c5'\n",
     "framework = 'PYTORCH'\n",
-    "framework_version = '0.4.0'\n",
+    "framework_version = '1.2.0'\n",
     "compiled_model_path = 's3://{}/{}/output'.format(bucket, compilation_job_name)"
    ]
   },