fix scheduler type definition

ourownstory · Aug 16, 2023 · 900eda1 · 900eda1
1 parent 6d137ae
commit 900eda1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neuralprophet/configure.py b/neuralprophet/configure.py
@@ -94,7 +94,7 @@ class Train:
     optimizer: Union[str, Type[torch.optim.Optimizer]]
     quantiles: List[float] = field(default_factory=list)
     optimizer_args: dict = field(default_factory=dict)
-    scheduler: Optional[Type[torch.optim.lr_scheduler._LRScheduler]] = None
+    scheduler: Optional[Type[torch.optim.lr_scheduler.OneCycleLR]] = None
     scheduler_args: dict = field(default_factory=dict)
     newer_samples_weight: float = 1.0
     newer_samples_start: float = 0.0