chore: fix QuantizationAwarweTraining notebook

ref #760
zama-ai · Jun 27, 2024 · 27954dd · 27954dd
1 parent 829b68b
commit 27954dd
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/docs/advanced_examples/QuantizationAwareTraining.ipynb b/docs/advanced_examples/QuantizationAwareTraining.ipynb
@@ -552,6 +552,16 @@
     "torch_model.unprune()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Move torch_model to CPU\n",
+    "torch_model = torch_model.cpu()"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 11,