more layers

app.py

@@ -16,9 +16,21 @@ def predict(img):
     topk_values, topk_indices = torch.topk(output, 2)  # Get the top 2 classes
     return [str(k) for k in topk_indices[0].tolist()]
 
+with gr.Blocks() as iface:
+    gr.Markdown("# MNIST + Gradio End to End")
+    gr.HTML("Shows end to end MNIST training with Gradio interface")
+    with gr.Row():
+        with gr.Column():
+            sp = gr.Sketchpad(shape=(28, 28))
+            with gr.Row():
+                with gr.Column():
+                    pred_button = gr.Button("Predict")
+                with gr.Column():
+                    clear = gr.Button("Clear")
+        with gr.Column():
+            label1 = gr.Label(label='1st Pred')
+            label2 = gr.Label(label='2nd Pred')
 
-sp = gr.Sketchpad(shape=(28, 28))
-
-gr.Interface(fn=predict,
-             inputs=sp,
-             outputs=['label','label']).launch()
+    pred_button.click(predict, inputs=sp, outputs=[label1,label2])
+    clear.click(lambda: None, None, sp, queue=False)
+iface.launch()

model.py

@@ -5,11 +5,13 @@ class Net(nn.Module):
     def __init__(self):
         super(Net, self).__init__()
         self.fc1 = nn.Linear(28*28, 128)  # MNIST images are 28x28
-        self.fc2 = nn.Linear(128, 64)
-        self.fc3 = nn.Linear(64, 10)  # There are 10 classes (0 through 9)
+        self.fc2 = nn.Linear(128, 128)
+        self.fc3 = nn.Linear(128, 64)
+        self.fc4 = nn.Linear(64, 10)  # There are 10 classes (0 through 9)
 
     def forward(self, x):
         x = x.view(x.shape[0], -1)  # Flatten the input
         x = torch.relu(self.fc1(x))
         x = torch.relu(self.fc2(x))
-        return self.fc3(x)
+        x = torch.relu(self.fc3(x))
+        return self.fc4(x)

requirements.txt

@@ -1,5 +1,3 @@
-gradio==3.29.0
-numpy==1.23.5
 Pillow==9.1.0
 torch==2.0.1
 torchvision==0.15.2