Update src/model.py

src/model.py

@@ -20,7 +20,9 @@ class Model_Load:
           # best_model_path='models/consumer_final_10/lightning_logs/lightning_logs/version_0/checkpoints/epoch=5-step=49260.ckpt'
           # best_tft = TemporalFusionTransformer.load_from_checkpoint(best_model_path)
           filename="models/cpu_energy_tft_model_v1.sav"
-          best_tft=pickle.load(open(filename, 'rb'))
+          # best_tft=pickle.load(open(filename, 'rb'))
+          with open(filename, 'rb') as f:  
+              best_tft=pd.read_pickle(f)
           print('Model Load Sucessfully.')
           return best_tft
         elif model_option=='Prophet':
@@ -42,7 +44,9 @@ class Model_Load:
           # best_model_path="models/store_item_10_lead_1_v2/lightning_logs/lightning_logs/version_2/checkpoints/epoch=7-step=4472.ckpt"
           # best_model_path="models/store_item_10_lead_1_v3/lightning_logs/lightning_logs/version_0/checkpoints/epoch=7-step=4472.ckpt"
           # best_tft = TemporalFusionTransformer.load_from_checkpoint(best_model_path)
-          best_tft=pickle.load(open("models/cpu_finalized_model_v1.sav", 'rb'))
+          # best_tft=pickle.load(open("models/cpu_finalized_model_v1.sav", 'rb'))
+          with open("models/cpu_finalized_model_v1.sav", 'rb') as f:  
+              best_tft=pd.read_pickle(f)
           print('Model Load Sucessfully.')
           return best_tft
         elif model_option=='Prophet':