Upload app.py

app.py

@@ -45,16 +45,7 @@ def fn(glycan, model):
     glycan = [glycan]
     label = [0]
     data = next(iter(dataset_to_dataloader(glycan, label, batch_size=1)))
-    device = "cpu"
-    if torch.cuda.is_available():
-      device = "cuda:0"
-    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(x,edge_index, batch).cpu().detach().numpy()[0]
+    pred = model(data)[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)}