bug fix

app.py

@@ -59,7 +59,7 @@ def predict(_model, _dataloader, datepicker):
     preds = raw_preds_to_df(out, quantiles = None)
     return preds[["pred_idx", "Group", "pred"]]
 
-def update_plot(df, preds, axs):
+def update_plot(df, preds, axs, fig):
     df = pd.merge(df, preds, left_on=["time_idx", "Group"], right_on=["pred_idx", "Group"], how = "left")
     df = df[~df["pred"].isna()]
     df[["sales", "pred"]] = df[["sales", "pred"]].replace(0.0, np.nan)
@@ -135,7 +135,6 @@ def main():
     
         ### Experiments
     """)
-    st.write(df.head(5))
     rain = st.radio("Rain Indicator", ('Default', 'Yes', 'No'))
     
     temperature = st.slider('Change in Temperature', min_value=-10.0, max_value=10.0, value=0.0, step=0.25)
@@ -149,7 +148,7 @@ def main():
     if st.button("Forecast Sales", type="primary"):
         dataloader = prepare_dataset(parameters, df.copy(), rain, temperature, datepicker, rain_mapping)
         preds = predict(model, dataloader, datepicker)
-        update_plot(df, preds, axs)
+        update_plot(df, preds, axs, fig)
 
 if __name__ == '__main__':
     main()