Update app.py

app.py

@@ -264,7 +264,8 @@ def transfer_learning_forecasting():
             forecast_results['LSTM'] = generate_forecast(lstm_model, df)
         elif model_choice == "TFT":
             forecast_results['TFT'] = generate_forecast(tft_model, df)
-            
+
+        st.session_state.forecast_results = forecast_results
         for model_name, forecast_df in forecast_results.items():
             plot_forecasts(forecast_df.iloc[:horizon,:], df, f'{model_name} Forecast for {y_col}')
 
@@ -272,8 +273,6 @@ def transfer_learning_forecasting():
         time_taken = end_time - start_time
         st.success(f"Time taken for {model_choice} forecast: {time_taken:.2f} seconds")
 
-        st.session_state.forecast_results = forecast_results
-
     if 'forecast_results' in st.session_state:
         forecast_results = st.session_state.forecast_results
 
@@ -291,12 +290,13 @@ def transfer_learning_forecasting():
             #             )
     
         with tab_forecast:
-            df_grid = forecast_results[model_choice]
-            st.write(df_grid)
-            # grid_table = AgGrid(
-            #                 df_grid,
-            #                 theme="alpine",
-            #             )
+            if model_choice in forecast_results:
+                df_grid = forecast_results[model_choice]
+                st.write(df_grid)
+                # grid_table = AgGrid(
+                #                 df_grid,
+                #                 theme="alpine",
+                #             )
 
 
 def dynamic_forecasting():