diff --git a/Capsule Network.ipynb b/Capsule Network.ipynb
index f3cd682..91d6444 100644
--- a/Capsule Network.ipynb	
+++ b/Capsule Network.ipynb	
@@ -306,7 +306,7 @@
     "        loss.backward()\n",
     "        optimizer.step()\n",
     "\n",
-    "        train_loss += loss.data[0]\n",
+    "        train_loss += loss.data\n",
     "        \n",
     "        if batch_id % 100 == 0:\n",
     "            print \"train accuracy:\", sum(np.argmax(masked.data.cpu().numpy(), 1) == \n",
@@ -327,7 +327,7 @@
     "        output, reconstructions, masked = capsule_net(data)\n",
     "        loss = capsule_net.loss(data, output, target, reconstructions)\n",
     "\n",
-    "        test_loss += loss.data[0]\n",
+    "        test_loss += loss.data\n",
     "        \n",
     "        if batch_id % 100 == 0:\n",
     "            print \"test accuracy:\", sum(np.argmax(masked.data.cpu().numpy(), 1) == \n",