Update app.py

app.py

@@ -6,7 +6,22 @@ from PIL import Image
 learn = torch.load('digit_classifier.pth')
 learn.eval() #switch to eval mode
 labels = [str(x) for x in range(10)]
+#################################
+#Define class for importing Model
+class DigitClassifier(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = torch.nn.Linear(64, 32)
+        self.fc2 = torch.nn.Linear(32, 16)
+        self.fc3 = torch.nn.Linear(16, 10)
 
+    def forward(self, x):
+        x = x.view(-1, 64)
+        x = torch.relu(self.fc1(x))
+        x = torch.relu(self.fc2(x))
+        x = self.fc3(x)
+        return x
+#########################################
 #Define function to reduce image of arbitrary size to 8x8 per model requirements.
 def reduce_image_count(image):
     output_size = (8, 8)
@@ -18,9 +33,11 @@ def reduce_image_count(image):
             block = image[i*block_size[0]:(i+1)*block_size[0], j*block_size[1]:(j+1)*block_size[1]]
             count = np.count_nonzero(block)
             output[i, j] = 16 - ((count / (block_size[0] * block_size[1])) * 16)
-            
+
     return output
 
+#########################################
+
 def predict(img):
     #First take input and reduce it to 8x8 px as the dataset was
     pil_image = Image.open(img) #get image