Update app.py

app.py

@@ -54,7 +54,7 @@ def load_all_models():
     return nhits_models, timesnet_models, lstm_models, tft_models
 
 def generate_forecast(model, df,tag=False):
-    if tag:
+    if tag == 'retrain':
         forecast_df = model.predict()
     else:
         forecast_df = model.predict(df=df)
@@ -162,7 +162,7 @@ def model_train(df,model, ds_col, freq):
     nf = NeuralForecast(models=[model], freq=freq)
     df[ds_col] = pd.to_datetime(df[ds_col])
     nf.fit(df)
-    return model
+    return nf
 
 def forecast_time_series(df, model_type, horizon, max_steps=200, ds_col='ds',y_col):
     start_time = time.time()  # Start timing