Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import pandas as pd
 from neuralprophet import NeuralProphet
 import warnings
-from torch.optim.lr_scheduler import OneCycleLR
+import torch.optim as optim
+from torch.optim.lr_scheduler import LambdaLR
 
 warnings.filterwarnings("ignore", category=UserWarning)
 
@@ -20,12 +21,9 @@ class CustomNeuralProphet(NeuralProphet):
         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:
-            if "lr_scheduler" in param_group:
-                lr_scheduler = param_group["lr_scheduler"]
-                lr_scheduler.step()
-
+        # Custom logic for LR scheduler step
+        for lr_scheduler in optimizer.param_groups[0]['lr_scheduler']:
+            lr_scheduler.step()
 
 m = CustomNeuralProphet(
     n_forecasts=30,
@@ -44,8 +42,12 @@ m = CustomNeuralProphet(
     learning_rate=0.03,
 )
 
+# Set the custom LR scheduler
+m.fit(df, freq='D')  # Fit the model first before accessing the optimizer
+m.optimizer = optim.Adam(m.model.parameters(), lr=0.03)  # Example optimizer, adjust as needed
 
-m.fit(df, freq='D')
+lr_scheduler = LambdaLR(m.optimizer, lambda epoch: 0.95 ** epoch)  # Example LR scheduler, adjust as needed
+m.optimizer.param_groups[0]['lr_scheduler'] = [lr_scheduler]
 
 future = m.make_future_dataframe(df, periods=30, n_historic_predictions=True)
 forecast = m.predict(future)
@@ -72,3 +74,4 @@ if __name__ == "__main__":
 
 
 
+