diff --git a/notebooks/example.ipynb b/notebooks/example.ipynb
index 49741a6..faf52c4 100644
--- a/notebooks/example.ipynb
+++ b/notebooks/example.ipynb
@@ -174,16 +174,16 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Epoch 0: 0.01 Kappa Score^2\n",
+      "Epoch 0: 0.03 Kappa Score^2\n",
       "Epoch 5: 0.01 Kappa Score^2\n",
-      "Epoch 10: 0.02 Kappa Score^2\n",
+      "Epoch 10: 0.01 Kappa Score^2\n",
       "Epoch 15: 0.03 Kappa Score^2\n",
-      "Epoch 20: 0.04 Kappa Score^2\n",
-      "Epoch 25: 0.06 Kappa Score^2\n",
-      "Epoch 30: 0.08 Kappa Score^2\n",
-      "Epoch 35: 0.09 Kappa Score^2\n",
-      "Epoch 40: 0.10 Kappa Score^2\n",
-      "Epoch 45: 0.12 Kappa Score^2\n"
+      "Epoch 20: 0.05 Kappa Score^2\n",
+      "Epoch 25: 0.07 Kappa Score^2\n",
+      "Epoch 30: 0.09 Kappa Score^2\n",
+      "Epoch 35: 0.10 Kappa Score^2\n",
+      "Epoch 40: 0.12 Kappa Score^2\n",
+      "Epoch 45: 0.13 Kappa Score^2\n"
      ]
     }
    ],
@@ -212,7 +212,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Test Kappa Score^2: 0.24\n"
+      "Test Kappa Score^2: 0.19\n"
      ]
     }
    ],
@@ -238,7 +238,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -247,7 +247,7 @@
     "optim = torch.optim.AdamW(model.parameters(), lr=1e-3)\n",
     "loss_cdw = CDW_CELoss(num_classes=7,\n",
     "                  transform=\"huber\",\n",
-    "                  alpha=2.).to(device)\n",
+    "                  alpha=1.5).to(device)\n",
     "loss = nn.CrossEntropyLoss().to(device)\n",
     "\n",
     "\n",
@@ -260,73 +260,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "tensor([2, 4, 4, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "Epoch 0: 0.05 Kappa Score^2\n",
-      "tensor([2, 3, 3, 2]) tensor([2, 2, 2, 2], device='cuda:0')\n",
-      "tensor([2, 3, 2, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 3, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([1, 4, 2, 2]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([4, 2, 4, 5]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "Epoch 5: 0.06 Kappa Score^2\n",
-      "tensor([2, 2, 2, 3]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([5, 3, 4, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 3, 2, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 3, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 2, 2]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "Epoch 10: 0.08 Kappa Score^2\n",
-      "tensor([3, 3, 4, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([4, 4, 3, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([2, 3, 4, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([2, 3, 2, 2]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 2, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
+      "Epoch 0: 0.00 Kappa Score^2\n",
+      "Epoch 5: 0.05 Kappa Score^2\n",
+      "Epoch 10: 0.07 Kappa Score^2\n",
       "Epoch 15: 0.09 Kappa Score^2\n",
-      "tensor([3, 3, 2, 2]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([3, 3, 4, 2]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "tensor([4, 3, 3, 2]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([1, 4, 3, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([2, 3, 3, 3]) tensor([2, 3, 3, 3], device='cuda:0')\n",
       "Epoch 20: 0.11 Kappa Score^2\n",
-      "tensor([3, 3, 3, 1]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 3, 2, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 3, 2]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([4, 3, 2, 2]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 3, 3, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "Epoch 25: 0.12 Kappa Score^2\n",
-      "tensor([2, 2, 3, 2]) tensor([3, 2, 2, 3], device='cuda:0')\n",
-      "tensor([4, 4, 3, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 3, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([5, 4, 3, 3]) tensor([3, 3, 4, 3], device='cuda:0')\n",
-      "tensor([2, 2, 3, 3]) tensor([2, 2, 2, 3], device='cuda:0')\n",
-      "Epoch 30: 0.13 Kappa Score^2\n",
-      "tensor([4, 3, 5, 4]) tensor([3, 3, 3, 4], device='cuda:0')\n",
-      "tensor([4, 2, 3, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 3, 3, 2]) tensor([3, 3, 4, 2], device='cuda:0')\n",
-      "tensor([3, 2, 1, 2]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([4, 3, 2, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "Epoch 35: 0.15 Kappa Score^2\n",
-      "tensor([2, 2, 3, 3]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "tensor([4, 3, 3, 4]) tensor([3, 3, 2, 3], device='cuda:0')\n",
-      "tensor([4, 4, 3, 3]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 3, 3, 3]) tensor([2, 3, 3, 3], device='cuda:0')\n",
-      "tensor([4, 3, 4, 3]) tensor([3, 3, 4, 3], device='cuda:0')\n",
-      "Epoch 40: 0.16 Kappa Score^2\n",
-      "tensor([3, 4, 1, 2]) tensor([3, 3, 2, 2], device='cuda:0')\n",
-      "tensor([2, 3, 3, 2]) tensor([2, 2, 3, 3], device='cuda:0')\n",
-      "tensor([4, 2, 3, 4]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "tensor([2, 2, 3, 2]) tensor([3, 3, 3, 3], device='cuda:0')\n",
-      "tensor([3, 2, 4, 2]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "Epoch 45: 0.17 Kappa Score^2\n",
-      "tensor([2, 2, 4, 2]) tensor([2, 2, 3, 3], device='cuda:0')\n",
-      "tensor([3, 4, 2, 3]) tensor([3, 4, 3, 3], device='cuda:0')\n",
-      "tensor([2, 3, 2, 2]) tensor([3, 2, 3, 3], device='cuda:0')\n",
-      "tensor([4, 4, 3, 4]) tensor([3, 3, 3, 3], device='cuda:0')\n"
+      "Epoch 25: 0.13 Kappa Score^2\n",
+      "Epoch 30: 0.14 Kappa Score^2\n",
+      "Epoch 35: 0.16 Kappa Score^2\n",
+      "Epoch 40: 0.18 Kappa Score^2\n",
+      "Epoch 45: 0.19 Kappa Score^2\n"
      ]
     }
    ],
@@ -339,7 +289,7 @@
     "        probs = F.softmax(y_hat, dim=1)\n",
     "        pred = torch.argmax(y_hat, dim=1)\n",
     "        kappa_train.update(pred, y.to(device))\n",
-    "        l = 0.1 * loss_cdw(probs, y.to(device)) + 0.9 * loss(y_hat, y.to(device))\n",
+    "        l = 0.2 * loss_cdw(probs, y.to(device)) + 0.8 * loss(y_hat, y.to(device))\n",
     "        l.backward()\n",
     "        optim.step()\n",
     "        optim.zero_grad()\n",
@@ -349,7 +299,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {