Update app.py

app.py

@@ -1,19 +1,22 @@
 import gradio as gr
 from huggingface_hub import hf_hub_download
 import torch
+from clinicadl.utils.network.cnn.models import Conv5_FC3
 
 # Download model from Hub
 model_path = hf_hub_download(repo_id="ARAMIS-LAB/CNN-AD-CN", filename="model.pth.tar")
 
 # Load ClinicaDL model
-model = torch.load(model_path, map_location="cpu")
-print(model)
+checkpoint_state = torch.load(model_path, map_location="cpu")
+model = Conv5_FC3()
+model.load_state_dict(checkpoint_state["model"])
 model.eval()
 
 def predict(input_image):
     with torch.no_grad():
         output = model(input_image.unsqueeze(0))
-    return output.numpy().tolist()
+    output = output.squeeze(0).cpu().float()
+    return output
 
 demo = gr.Interface(fn=predict, inputs="image", outputs="label")
 demo.launch()