iface fixed

app.py

@@ -67,8 +67,8 @@ class ViT(nn.Module):
 # 加载模型
 model = ViT(num_classes=10)  # 确保num_classes与你的MNIST任务一致
 model_path = "vit_model.pth"  # 模型权重文件名
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-model.eval()
+model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'), weights_only=True))
+
 
 # 定义图像预处理
 transform = transforms.Compose([
@@ -106,8 +106,7 @@ def classify_image(image):
 
 iface = gr.Interface(
     fn=classify_image,
-    inputs=gr.Image(
-        source="canvas",
+    inputs=gr.Sketchpad(
         tool="editor",
         type="pil",
         invert_colors=True,
@@ -116,7 +115,7 @@ iface = gr.Interface(
         label="Draw a digit"
     ),
     outputs=gr.Label(num_top_classes=1),
-    title="MNIST Digit Classification with ViT",
+    title="MNIST Digit Classification with ViT", 
     description="使用鼠标手绘一个数字，模型将预测其所属的类别。"
 )