Upload app.py

app.py

@@ -34,18 +34,29 @@ model3 = torch.load("model3.pt", map_location=torch.device('cpu'))
 
 def fn(glycan, model):
     if model == "No data augmentation":
-      model = model1
-      model.eval()
+      model_pred = model1
+      model_pred.eval()
     elif model == "Ensemble":
-      model = model3
-      model.eval()
+      model_pred = model3
+      model_pred.eval()
     else:
-      model = model2
-      model.eval()
+      model_pred = model2
+      model_pred.eval()
+    
     glycan = [glycan]
     label = [0]
     data = next(iter(dataset_to_dataloader(glycan, label, batch_size=1)))
-    pred = model(data)
+    if model == "Ensemble":
+        pred = model_pred(data)
+    else:
+        device = "cpu"
+        x = data.labels
+        edge_index = data.edge_index
+        batch = data.batch
+        x = x.to(device)
+        edge_index = edge_index.to(device)
+        batch = batch.to(device)
+        pred = model_pred(x,edge_index, batch).cpu().detach().numpy()[0]
     pred = np.exp(pred)/sum(np.exp(pred)) # Softmax 
     pred = [float(x) for x in pred]
     pred = {class_list[i]:pred[i] for i in range(15)}