fixes per code review

sony · Jul 2, 2024 · 02f5c93 · 02f5c93
1 parent 429bb92
commit 02f5c93
Showing 1 changed file with 16 additions and 8 deletions.
diff --git a/tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_mobilenet_gptq.ipynb b/tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_mobilenet_gptq.ipynb
@@ -45,10 +45,23 @@
    "outputs": [],
    "source": [
     "!pip install -q torch torchvision\n",
-    "!pip install -q model-compression-toolkit\n",
     "!pip install -q tqdm"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "import importlib\n",
+    "if not importlib.util.find_spec('model_compression_toolkit'):\n",
+    "    !pip install model_compression_toolkit"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "6db5eec3f181f263"
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -210,16 +223,11 @@
    },
    "outputs": [],
    "source": [
-    "from model_compression_toolkit.gptq.common.gptq_constants import REG_DEFAULT\n",
-    "\n",
     "# Specify the target platform capability (TPC)\n",
     "tpc = mct.get_target_platform_capabilities(\"pytorch\", 'imx500', target_platform_version='v1')\n",
     "\n",
     "# Create a GPTQ quantization configuration and set the number of training iterations. \n",
-    "# For the sake of running faster, the hessian based weights are disabled in this tutorial.\n",
-    "gptq_config = mct.gptq.get_pytorch_gptq_config(n_epochs=50,\n",
-    "                                               use_hessian_based_weights=False,\n",
-    "                                               regularization_factor=REG_DEFAULT)"
+    "gptq_config = mct.gptq.get_pytorch_gptq_config(n_epochs=50)\n"
    ]
   },
   {
@@ -230,7 +238,7 @@
    },
    "source": [
     "### Run model Gradient-based Post-Training Quantization\n",
-    "Finally, we quantize our model using MCT's GPTQ API."
+    "Finally, we quantize our model using MCT's GPTQ API (this may take several minutes)."
    ]
   },
   {