bugfix

.ipynb_checkpoints/app-checkpoint.py

@@ -54,12 +54,12 @@ def prepare_dataset(parameters, df, rain, temperature, datepicker):
     return df.to_dataloader(train=False, batch_size=256,num_workers = 0)
 
 def predict(model, dataloader): 
-    return model.cpu().predict(dataloader, mode="raw", return_x=True, return_index=True)
+    return model.predict(dataloader, mode="raw", return_x=True, return_index=True)
     
 ## Initiate Data
 with open('data/parameters.pkl', 'rb') as f:
     parameters = pickle.load(f)
-model = TemporalFusionTransformer.cpu().load_from_checkpoint('model/tft_check.ckpt')
+model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
 
 df = pd.read_pickle('data/test_data.pkl')
 df = df.loc[(df["Branch"] == 15) & (df["Group"].isin(["6","7","4","1"]))]

app.py

@@ -54,12 +54,12 @@ def prepare_dataset(parameters, df, rain, temperature, datepicker):
     return df.to_dataloader(train=False, batch_size=256,num_workers = 0)
 
 def predict(model, dataloader): 
-    return model.cpu().predict(dataloader, mode="raw", return_x=True, return_index=True)
+    return model.predict(dataloader, mode="raw", return_x=True, return_index=True)
     
 ## Initiate Data
 with open('data/parameters.pkl', 'rb') as f:
     parameters = pickle.load(f)
-model = TemporalFusionTransformer.cpu().load_from_checkpoint('model/tft_check.ckpt')
+model = TemporalFusionTransformer.load_from_checkpoint('model/tft_check.ckpt', map_location=torch.device('cpu'))
 
 df = pd.read_pickle('data/test_data.pkl')
 df = df.loc[(df["Branch"] == 15) & (df["Group"].isin(["6","7","4","1"]))]