Update app.py

app.py

@@ -46,13 +46,14 @@ def predict():
     
     return features_dict
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 #model_link = snapshot_download(repo_id="FFZG-cleopatra/Croatian-News-Classifier")
 model_link = hf_hub_download(repo_id="FFZG-cleopatra/Croatian-News-Classifier",filename = "pytorch_model.bin")
 
 # multitask_model.from_pretrained(, config="/media/gaurish/angela/projects/CroatianSlovenEnglishBert/i-got-u-brother-cleopatra-workshop/src/models/multitask_model_3ep/config.json")
-multitask_model.load_state_dict(torch.load(model_link, map_location=torch.device('cpu')))
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+multitask_model.load_state_dict(torch.load(model_link, map_location=device))
+
 # multitask_model.to(device)
 predictions = []
 features_dict = predict()