fix: add fix for tutorial 07, leftover from #484

jaxleyverse · Nov 4, 2024 · 6509191 · 6509191
1 parent 345bb4b
commit 6509191
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/docs/tutorials/07_gradient_descent.ipynb b/docs/tutorials/07_gradient_descent.ipynb
@@ -823,8 +823,7 @@
     "for epoch in range(10):\n",
     "    epoch_loss = 0.0\n",
     "    for batch_ind, batch in enumerate(dataloader):\n",
-    "        current_batch = batch[0].numpy()\n",
-    "        label_batch = batch[1].numpy()\n",
+    "        current_batch, label_batch = batch[0]\n",
     "        loss_val, gradient = jitted_grad(opt_params, current_batch, label_batch)\n",
     "        updates, opt_state = optimizer.update(gradient, opt_state)\n",
     "        opt_params = optax.apply_updates(opt_params, updates)\n",