diff --git a/jupyter/agp_schedule.ipynb b/jupyter/agp_schedule.ipynb
index cce0734acb0c4ddb37728efb24a0e7bcfd25662b..c56560730655a5b36e4e6b74cdae9ff04d7781c7 100644
--- a/jupyter/agp_schedule.ipynb
+++ b/jupyter/agp_schedule.ipynb
@@ -116,7 +116,7 @@
     "            \n",
     "            sparsity_levels.append(sparsity)\n",
     "            sparsity.backward()\n",
-    "            d_sparsity.append(current_epoch.grad.data[0])\n",
+    "            d_sparsity.append(current_epoch.grad.item())\n",
     "            current_epoch.grad.data.zero_()\n",
     "        else:\n",
     "            sparsity_levels.append(sparsity)\n",