Update src/nf.py

src/nf.py

@@ -136,11 +136,7 @@ def forecast_pretrained_model(Y_df: pd.DataFrame, model: str, fh: int, max_steps
 
     # Load the checkpoint and initialize NHITS with required parameters
     file_ = f"./models/{model}.ckpt"
-    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'])
+    nhits = NeuralForecast.load_from_checkpoint(file_)
 
     # Fit
     if max_steps > 0: