Update pretrained_weights.ipynb

Fixed an error in the state dict loading of the turorial and added a comment on the num_classes parameter when creating timm models.

docs/tutorials/pretrained_weights.ipynb

@@ -227,8 +227,26 @@
    "outputs": [],
    "source": [
     "in_chans = weights.meta[\"in_chans\"]\n",
-    "model = timm.create_model(\"resnet18\", in_chans=in_chans, num_classes=10)\n",
-    "model = model.load_state_dict(weights.get_state_dict(progress=True), strict=False)"
+    "model = timm.create_model(\"resnet18\", in_chans=in_chans, num_classes=0)\n",
+    "model.load_state_dict(weights.get_state_dict(progress=True), strict=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Setting `num_classes=0` will prevent the creation of a prediction head (fully-connected layer). However, you can create a timm model with a prediction head and still match the keys of all but the last fully-connected layer, the parameters of which will be randomly initialized, when using the `strict=False` flag."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "in_chans = weights.meta[\"in_chans\"]\n",
+    "model = timm.create_model(\"resnet18\", in_chans=in_chans, num_classes=256)\n",
+    "model.load_state_dict(weights.get_state_dict(progress=True), strict=False)"
    ]
   },
   {