diff --git a/docs/source/tutorial.ipynb b/docs/source/tutorial.ipynb
index 7f39c09..2ecfffa 100644
--- a/docs/source/tutorial.ipynb
+++ b/docs/source/tutorial.ipynb
@@ -2,223 +2,252 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "7ac5eff0",
+   "metadata": {},
    "source": [
     "# A short tutorial for the `prototorch.models` plugin"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "beb83780",
+   "metadata": {},
    "source": [
     "## Introduction"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "43b74278",
+   "metadata": {},
    "source": [
     "This is a short tutorial for the [models](https://github.com/si-cim/prototorch_models) plugin of the [ProtoTorch](https://github.com/si-cim/prototorch) framework.\n",
     "\n",
     "[ProtoTorch](https://github.com/si-cim/prototorch) provides [torch.nn](https://pytorch.org/docs/stable/nn.html) modules and utilities to implement prototype-based models. However, it is up to the user to put these modules together into models and handle the training of these models. Expert machine-learning practioners and researchers sometimes prefer this level of control. However, this leads to a lot of boilerplate code that is essentially same across many projects. Needless to say, this is a source of a lot of frustration. [PyTorch-Lightning](https://pytorch-lightning.readthedocs.io/en/latest/) is a framework that helps avoid a lot of this frustration by handling the boilerplate code for you so you don't have to reinvent the wheel every time you need to implement a new model.\n",
     "\n",
     "With the [prototorch.models](https://github.com/si-cim/prototorch_models) plugin, we've gone one step further and pre-packaged commonly used prototype-models like GMLVQ as [Lightning-Modules](https://pytorch-lightning.readthedocs.io/en/latest/api/pytorch_lightning.core.lightning.html?highlight=lightning%20module#pytorch_lightning.core.lightning.LightningModule). With only a few lines to code, it is now possible to build and train prototype-models. It quite simply cannot get any simpler than this."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "4e5d1fad",
+   "metadata": {},
    "source": [
     "## Basics"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "1244b66b",
+   "metadata": {},
    "source": [
     "First things first. When working with the models plugin, you'll probably need `torch`, `prototorch` and `pytorch_lightning`. So, we recommend that you import all three like so:"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "dcb88e8a",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import prototorch as pt\n",
     "import pytorch_lightning as pl\n",
     "import torch"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "1adbe2f8",
+   "metadata": {},
    "source": [
     "### Building Models"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "96663ab1",
+   "metadata": {},
    "source": [
     "Let's start by building a `GLVQ` model. It is one of the simplest models to build. The only requirements are a prototype distribution and an initializer."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "819ba756",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "model = pt.models.GLVQ(\n",
     "    hparams=dict(distribution=[1, 1, 1]),\n",
     "    prototypes_initializer=pt.initializers.ZerosCompInitializer(2),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1b37e97c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "print(model)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "d2c86903",
+   "metadata": {},
    "source": [
     "The key `distribution` in the `hparams` argument describes the prototype distribution. If it is a Python [list](https://docs.python.org/3/tutorial/datastructures.html), it is assumed that there are as many entries in this list as there are classes, and the number at each location of this list describes the number of prototypes to be used for that particular class. So, `[1, 1, 1]` implies that we have three classes with one prototype per class.   If it is a Python [tuple](https://docs.python.org/3/tutorial/datastructures.html), a shorthand of `(num_classes, prototypes_per_class)` is assumed. If it is a Python [dictionary](https://docs.python.org/3/tutorial/datastructures.html), the key-value pairs describe the class label and the number of prototypes for that class respectively. So, `{0: 2, 1: 2, 2: 2}` implies that we have three classes with labels `{1, 2, 3}`, each equipped with two prototypes. If however, the dictionary contains the keys `\"num_classes\"` and `\"per_class\"`, they are parsed to use their values as one might expect.\n",
     "\n",
     "The `prototypes_initializer` argument describes how the prototypes are meant to be initialized. This argument has to be an instantiated object of some kind of [AbstractComponentsInitializer](https://github.com/si-cim/prototorch/blob/dev/prototorch/components/initializers.py#L18). If this is a [ShapeAwareCompInitializer](https://github.com/si-cim/prototorch/blob/dev/prototorch/components/initializers.py#L41), this only requires a `shape` arugment that describes the shape of the prototypes. So, `pt.initializers.ZerosCompInitializer(3)` creates 3d-vector prototypes all initialized to zeros."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "45806052",
+   "metadata": {},
    "source": [
     "### Data"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "9d62c4c6",
+   "metadata": {},
    "source": [
     "The preferred way to working with data in `torch` is to use the [Dataset and Dataloader API](https://pytorch.org/tutorials/beginner/basics/data_tutorial.html). There a few pre-packaged datasets available under `prototorch.datasets`. See [here](https://prototorch.readthedocs.io/en/latest/api.html#module-prototorch.datasets) for a full list of available datasets."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "504df02c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = pt.datasets.Iris(dims=[0, 2])"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "3b8e7756",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "type(train_ds)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "bce43afa",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds.data.shape, train_ds.targets.shape"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "26a83328",
+   "metadata": {},
    "source": [
     "Once we have such a dataset, we could wrap it in a `Dataloader` to load the data in batches, and possibly apply some transformations on the fly."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "67b80fbe",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=2)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "c1185f31",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "type(train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "9b5a8963",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "x_batch, y_batch = next(iter(train_loader))\n",
     "print(f\"{x_batch=}, {y_batch=}\")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "dd492ee2",
+   "metadata": {},
    "source": [
     "This perhaps seems like a lot of work for a small dataset that fits completely in memory. However, this comes in very handy when dealing with huge datasets that can only be processed in batches."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "5176b055",
+   "metadata": {},
    "source": [
     "### Training"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "46a7a506",
+   "metadata": {},
    "source": [
     "If you're familiar with other deep learning frameworks, you might perhaps expect a `.fit(...)` or `.train(...)` method. However, in PyTorch-Lightning, this is done slightly differently. We first create a trainer and then pass the model and the Dataloader to `trainer.fit(...)` instead. So, it is more functional in style than object-oriented."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "279e75b7",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "trainer = pl.Trainer(max_epochs=2, weights_summary=None)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "e496b492",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "trainer.fit(model, train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "497fbff6",
+   "metadata": {},
    "source": [
     "### From data to a trained model - a very minimal example"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "ab069c5d",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = pt.datasets.Iris(dims=[0, 2])\n",
     "train_loader = torch.utils.data.DataLoader(train_ds, batch_size=32)\n",
@@ -230,49 +259,189 @@
     "\n",
     "trainer = pl.Trainer(max_epochs=50, weights_summary=None)\n",
     "trainer.fit(model, train_loader)"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "30c71a93",
+   "metadata": {},
    "source": [
-    "## Advanced"
-   ],
-   "metadata": {}
+    "### Saving/Loading trained models"
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "f74ed2c1",
+   "metadata": {},
    "source": [
-    "### Initializing prototypes with a subset of a dataset (along with transformations)"
-   ],
-   "metadata": {}
+    "Pytorch Lightning can automatically checkpoint the model during various stages of training, but it also possible to manually save a checkpoint after training."
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "3156658d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ckpt_path = \"./checkpoints/glvq_iris.ckpt\"\n",
+    "trainer.save_checkpoint(ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c1c34055",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "loaded_model = pt.models.GLVQ.load_from_checkpoint(ckpt_path, strict=False)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bbbb08e9",
+   "metadata": {},
+   "source": [
+    "### Visualizing decision boundaries in 2D"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "53ca52dc",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds).visualize(loaded_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8373531f",
+   "metadata": {},
+   "source": [
+    "### Saving/Loading trained weights"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "937bc458",
+   "metadata": {},
+   "source": [
+    "In most cases, the checkpointing workflow is sufficient. In some cases however, one might want to only save the trained weights from the model. The disadvantage of this method is that the model has be re-created using compatible initialization parameters before the weights could be loaded."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1f2035af",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ckpt_path = \"./checkpoints/glvq_iris_weights.pth\"\n",
+    "torch.save(model.state_dict(), ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1206021a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = pt.models.GLVQ(\n",
+    "    dict(distribution=(3, 2)),\n",
+    "    prototypes_initializer=pt.initializers.ZerosCompInitializer(2),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9f2a4beb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds, title=\"Before loading the weights\").visualize(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "528d2fc2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.load(ckpt_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ec817e6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.load_state_dict(torch.load(ckpt_path), strict=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a208eab7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pt.models.VisGLVQ2D(data=train_ds, title=\"After loading the weights\").visualize(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f8de748f",
+   "metadata": {},
+   "source": [
+    "## Advanced"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1f6a33a5",
+   "metadata": {},
+   "source": [
+    "### Initializing prototypes with a subset of a dataset (along with transformations)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "946ce341",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import prototorch as pt\n",
     "import pytorch_lightning as pl\n",
     "import torch\n",
     "from torchvision import transforms\n",
     "from torchvision.datasets import MNIST"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "510d9bd4",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "from matplotlib import pyplot as plt"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "ea7c1228",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "train_ds = MNIST(\n",
     "    \"~/datasets\",\n",
@@ -284,59 +453,64 @@
     "        transforms.ToTensor(),\n",
     "    ]),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1b9eaf5c",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "s = int(0.05 * len(train_ds))\n",
     "init_ds, rest_ds = torch.utils.data.random_split(train_ds, [s, len(train_ds) - s])"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "8c32c9f2",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "init_ds"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "68a9a8b9",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "model = pt.models.ImageGLVQ(\n",
     "    dict(distribution=(10, 5)),\n",
     "    prototypes_initializer=pt.initializers.SMCI(init_ds),\n",
     ")"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "6f23df86",
+   "metadata": {},
+   "outputs": [],
    "source": [
     "plt.imshow(model.get_prototype_grid(num_columns=10))"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "4fa69f92",
+   "metadata": {},
    "source": [
     "## FAQs"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "fa20f9ac",
+   "metadata": {},
    "source": [
     "### How do I Retrieve the prototypes and their respective labels from the model?\n",
     "\n",
@@ -351,11 +525,12 @@
     "```python\n",
     ">>> model.prototype_labels\n",
     "```"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "id": "ba8215bf",
+   "metadata": {},
    "source": [
     "### How do I make inferences/predictions/recall with my trained model?\n",
     "\n",
@@ -370,13 +545,12 @@
     "```python\n",
     ">>> y_pred = model(torch.Tensor(x_train))  # returns probabilities\n",
     "```"
-   ],
-   "metadata": {}
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -390,7 +564,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.4"
+   "version": "3.9.9"
   }
  },
  "nbformat": 4,