diff --git a/examples/parametrisations.ipynb b/examples/parametrisations.ipynb
index 06141c82..4cd3cee0 100644
--- a/examples/parametrisations.ipynb
+++ b/examples/parametrisations.ipynb
@@ -67,7 +67,7 @@
    "outputs": [],
    "source": [
     "class SymmetricParametrization(nn.Module):\n",
-    "    def forward(X):\n",
+    "    def forward(self, X):\n",
     "        A = X.triu()\n",
     "        return A + A.T "
    ]
@@ -113,9 +113,10 @@
    "source": [
     "class PruningParametrization(nn.Module):\n",
     "    def __init__(self, X, p_drop=0.2):\n",
+    "        super().__init__()\n",
     "        # sample zeros with probability p_drop\n",
     "        mask = torch.full_like(X, 1.0 - p_drop)\n",
-    "        self.mask = torch.bernoulli(mask)\n",
+    "        self.register_buffer(\"mask\", torch.bernoulli(mask))\n",
     "\n",
     "    def forward(self, X):\n",
     "        return X * self.mask"
@@ -313,8 +314,6 @@
     "        return torch.allclose(X, -X.T)\n",
     "\n",
     "    def right_inverse(self, X):\n",
-    "        if not self.is_skew(X):\n",
-    "            raise ValueError(\"This  matirx is not skew-symmetric!\")\n",
     "        return X.triu(1)\n",
     "    \n",
     "# Skew.forward(Skew.right_inverse(X)) == X\n",
@@ -347,11 +346,11 @@
     "    def forward(self, A):\n",
     "        # Cayley map: (I + A)(I - A)^{-1}\n",
     "        # This is orthogonal whenever A is skew-symmetric\n",
-    "        Id = torch.eye(A.size(0))\n",
+    "        Id = torch.eye(A.size(0), dtype=A.dtype, device=A.device)\n",
     "        return self.B @ torch.linalg.solve(Id - A, Id + A)\n",
     "\n",
     "    def is_orthogonal(self, X):\n",
-    "        Id = torch.eye(X.size(0))\n",
+    "        Id = torch.eye(X.size(0), dtype=X.dtype, device=X.device)\n",
     "        return torch.allclose(X.T @ X, Id, atol=1e-6)\n",
     "\n",
     "    def right_inverse(self, X):\n",
@@ -362,7 +361,7 @@
     "        return torch.zeros_like(X)\n",
     "\n",
     "\n",
-    "model = nn.Linear(5,5)\n",
+    "model = nn.Linear(5, 5)\n",
     "P.register_parametrization(model, \"weight\", Skew())\n",
     "P.register_parametrization(model, \"weight\", Orthogonal(5))\n",
     "\n",