Added map_location

app.py

@@ -20,7 +20,7 @@ D_LABITEMS = pd.read_csv("D_LABITEMS.csv", header = "infer", sep = ",", encoding
 
 def load_model():
     path = r"final_model.pt"
-    kwargs, state = torch.load(path, weights_only=False)
+    kwargs, state = torch.load(path, weights_only=False, map_location=device)
     model = VariationalGNN(**kwargs).to(device)
     model.load_state_dict(state)
     return model