Fix PyTorch lightning total steps

MilesCranmer · Apr 18, 2023 · 1e43bdc · 1e43bdc
1 parent 3f60d16
commit 1e43bdc
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/examples/pysr_demo.ipynb b/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",