Update app.py

app.py

@@ -3,6 +3,8 @@ import pandas as pd
 from neuralprophet import NeuralProphet
 import io
 import warnings
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.optim.optimizer import Optimizer
 
 warnings.filterwarnings("ignore", category=UserWarning)
 
@@ -13,13 +15,14 @@ df = df.rename(columns={"Date": "ds", "Price": "y"})
 df.fillna(method='ffill', inplace=True)
 df.dropna(inplace=True)
 
+
 class CustomNeuralProphet(NeuralProphet):
     def lr_scheduler_step(self, epoch, batch_idx, optimizer):
-        # Get the OneCycleLR scheduler
-        scheduler = self.optimizers[0].scheduler
-
-        # Call the `step` method on the scheduler
-        scheduler.step()
+        # 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()
 
 m = CustomNeuralProphet(
     n_forecasts=30,
@@ -38,6 +41,14 @@ m = CustomNeuralProphet(
     learning_rate=0.03,
 )
 
+# Set the custom LR scheduler
+optimizer = m.trainer.optimizers[0]
+optimizer.lr_scheduler = {
+    "scheduler": OneCycleLR(optimizer, max_lr=0.1, steps_per_epoch=100, epochs=10),
+    "interval": "step",
+    "frequency": 1,
+}
+
 m.fit(df, freq='D')
 
 future = m.make_future_dataframe(df, periods=30, n_historic_predictions=True)