Update src/nf.py

src/nf.py

@@ -12,6 +12,7 @@ from neuralforecast.core import NeuralForecast
 from neuralforecast.auto import NHITS as autoNHITS
 from neuralforecast.tsdataset import TimeSeriesDataset, TimeSeriesLoader
 from neuralforecast.models import NHITS
+import torch
 
 # GLOBAL PARAMETERS
 DEFAULT_HORIZON = 30
@@ -133,9 +134,13 @@ def forecast_pretrained_model(Y_df: pd.DataFrame, model: str, fh: int, max_steps
     scaler.fit(Y_df)
     Y_df = scaler.transform(Y_df)
 
-    # Model
+    # Load the checkpoint and initialize NHITS with required parameters
     file_ = f"./models/{model}.ckpt"
-    nhits = NHITS.load_from_checkpoint(file_)
+    checkpoint = torch.load(file_)
+    h = checkpoint['hyper_parameters']['h']
+    input_size = checkpoint['hyper_parameters']['input_size']
+    nhits = NHITS(h=h, input_size=input_size)
+    nhits.load_state_dict(checkpoint['state_dict'])
 
     # Fit
     if max_steps > 0: