Update app.py

app.py

@@ -8,8 +8,12 @@ model_path = hf_hub_download(repo_id="ARAMIS-LAB/CNN-AD-CN", filename="model.pth
 
 # Load ClinicaDL model
 checkpoint_state = torch.load(model_path, map_location="cpu")
-print(checkpoint_state)
-model = Conv5_FC3()
+model = Conv5_FC3(input_size= [
+        1,
+        169,
+        208,
+        179
+    ])
 model.load_state_dict(checkpoint_state["model"])
 model.eval()
 
@@ -17,7 +21,7 @@ def predict(input_image):
     with torch.no_grad():
         output = model(input_image.unsqueeze(0))
     output = output.squeeze(0).cpu().float()
-    return output
+    return output[0]
 
 demo = gr.Interface(fn=predict, inputs="image", outputs="label")
 demo.launch()