Merge pull request #295 from MilesCranmer/fix-colab

examples/pysr_demo.ipynb

@@ -1118,14 +1118,12 @@
         "        return self.training_step(batch, batch_idx)\n",
         "\n",
         "    def configure_optimizers(self):\n",
-        "        self.trainer.reset_train_dataloader()\n",
-        "\n",
         "        optimizer = torch.optim.Adam(self.parameters(), lr=self.max_lr)\n",
         "        scheduler = {\n",
         "            \"scheduler\": torch.optim.lr_scheduler.OneCycleLR(\n",
         "                optimizer,\n",
         "                max_lr=self.max_lr,\n",
-        "                total_steps=self.total_steps,\n",
+        "                total_steps=self.trainer.estimated_stepping_batches,\n",
         "                final_div_factor=1e4,\n",
         "            ),\n",
         "            \"interval\": \"step\",\n",