Upload 3 files

app.py

@@ -2,35 +2,9 @@ import torch
 from torch import nn
 import gradio as gr
 
-# Define the custom CNN model class that was trained on the MNIST data
-class CNN(nn.Module):
-    """
-    A custom CNN class.  The network has: (1) a convolution layer with 1 input channel and 16 output channels with ReLU activation and 2x2 max-pooling, (2) a second convolution layer with 16 input channels and 32 output channels with ReLU activation and 2x2 max-pooling, and (3) a linear output layer with 10 outputs.
-    """
-    def __init__(self):
-        super(CNN,self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(1,16,5,stride=1,padding=2),
-            nn.ReLU(),
-           	nn.MaxPool2d(kernel_size=2),
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(16,32,5,1,2),
-            nn.ReLU(),
-            nn.MaxPool2d(2),
-        )
-        self.out = nn.Linear(32*7*7,10)
-        
-    # Forward propogation method    
-    def forward(self,x):
-        x=self.conv1(x)
-        x=self.conv2(x)
-        x = x.view(-1,32*7*7)
-        return self.out(x)
-
-# Initialize an instance and load in the saved state_dict for the trained model
-model = CNN()
-model.load_state_dict(torch.load('mnist2.pkl',map_location=torch.device('cpu')))
+# Load the model and then the post-training state_dict
+model = torch.load('mnist_model.pth',map_location=torch.device('cpu'))
+model.load_state_dict(torch.load('mnist_model_weights.pth',map_location=torch.device('cpu')))
 model.eval()
 
 # Prediction function

mnist_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:170a7a2185e474af9dd12e1232e28a9ea6003451a49b6dbc225d98f89c072a0e
+size 13103247

mnist_model_weights.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4cd377c11e267b84a351cebf48d4ea0a768b05e5502e34a4bf2854f55bd7d19
+size 13100943