diff --git a/Capsule Network.ipynb b/Capsule Network.ipynb
index f3cd682..abe879a 100644
--- a/Capsule Network.ipynb	
+++ b/Capsule Network.ipynb	
@@ -170,7 +170,7 @@
     "        masked = Variable(torch.sparse.torch.eye(10))\n",
     "        if USE_CUDA:\n",
     "            masked = masked.cuda()\n",
-    "        masked = masked.index_select(dim=0, index=max_length_indices.squeeze(1).data)\n",
+    "        masked = masked.index_select(dim=0, index=max_length_indices.squeeze(1))\n",
     "        \n",
     "        reconstructions = self.reconstraction_layers((x * masked[:, :, None, None]).view(x.size(0), -1))\n",
     "        reconstructions = reconstructions.view(-1, 1, 28, 28)\n",