Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 import pandas as pd
 from neuralprophet import NeuralProphet
-import io
 import warnings
 from torch.optim.lr_scheduler import OneCycleLR
 
@@ -16,6 +15,10 @@ df.dropna(inplace=True)
 
 
 class CustomNeuralProphet(NeuralProphet):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.optimizer = None
+
     def lr_scheduler_step(self, epoch, batch_idx, optimizer):
         # Custom logic for OneCycleLR scheduler step
         for param_group in optimizer.param_groups:
@@ -23,8 +26,6 @@ class CustomNeuralProphet(NeuralProphet):
                 lr_scheduler = param_group["lr_scheduler"]
                 lr_scheduler.step()
 
-optimizer = None
-
 
 m = CustomNeuralProphet(
     n_forecasts=30,
@@ -44,9 +45,10 @@ m = CustomNeuralProphet(
 )
 
 # Set the custom LR scheduler
-optimizer = m.optimizer
-optimizer.lr_scheduler = {
-    "scheduler": OneCycleLR(optimizer),
+m.optimizer = m.model.optimizers[0]  # Use the optimizer from the model
+
+m.optimizer.lr_scheduler = {
+    "scheduler": OneCycleLR(m.optimizer),
     "interval": "step",
     "frequency": 1,
 }
@@ -77,3 +79,4 @@ if __name__ == "__main__":
 
 
 
+