Merge pull request #19 from kevin/main

(MNIST Tutorial) Fixing iterator usage resulting in slow training

examples/01_mnist.ipynb

@@ -578,7 +578,7 @@
     "def train(model, trainloader, testloader, iterations, test_every, device):\n",
     "\n",
     "  optimizer = torch.optim.AdamW(params=list(model.parameters()), lr=0.0001, eps=1e-8)\n",
-    "\n",
+    "  iterator = iter(trainloader)\n",
     "  model.train()\n",
     "  \n",
     "  train_losses = []\n",
@@ -595,7 +595,11 @@
     "      test_accuracy = None\n",
     "      for stepi in range(iterations):\n",
     "\n",
-    "          inputs, targets = next(iter(trainloader))\n",
+    "          try:\n",
+    "              inputs, targets = next(iterator)\n",
+    "          except StopIteration:\n",
+    "              iterator = iter(trainloader)\n",
+    "              inputs, targets = next(iterator)\n",
     "          inputs, targets = inputs.to(device), targets.to(device)\n",
     "          predictions, certainties, _ = model(inputs, track=False)\n",
     "          train_loss, where_most_certain = get_loss(predictions, certainties, targets)\n",