Upload with huggingface_hub

app.py

@@ -8,19 +8,23 @@ from dotenv import load_dotenv
 from torch import nn
 import torch.nn.functional as F
 
-class SimpleLenet(nn.Module):
+class Lenet(nn.Module):
     def __init__(self, args=None):
         super().__init__()
         self.conv1 = nn.Conv2d(1, 6, 5, padding=2) # -> 6 channels, 28x28
-        self.pool = nn.MaxPool2d(2) # -> 6 channels, 14x14
-        self.conv2 = nn.Conv2d(6, 120, 14) #-> 120 channels, 1x1
-        self.fc1 = nn.Linear(120, 10)
-        self.fc2 = nn.Linear(10, 10)
+        self.pool1 = nn.MaxPool2d(2) # -> 6 channels, 14x14
+        self.conv2 = nn.Conv2d(6, 16, 5) #-> 16 images, 10x10
+        self.pool2 = nn.MaxPool2d(2) # -> 16 channels,  5x5
+        self.conv3 = nn.Conv2d(16, 120, 5) #-> 16 images, 1x1
+        self.fc1 = nn.Linear(120, 84)
+        self.fc2 = nn.Linear(84, 10)
 
     def __call__(self, x):
         xx = F.relu(self.conv1(x))
-        xx = F.relu(self.pool(xx))
+        xx = F.relu(self.pool1(xx))
         xx = F.relu(self.conv2(xx))
+        xx = F.relu(self.pool2(xx))
+        xx = F.relu(self.conv3(xx))
         xx = xx.flatten(1)
         xx = F.relu(self.fc1(xx))
         return self.fc2(xx)
@@ -30,7 +34,7 @@ load_dotenv()
 hf_writer = gr.HuggingFaceDatasetSaver(os.getenv('HF_TOKEN'), "simple-mnist-flagging")
 
 def load_model():
-    model = SimpleLenet()
+    model = Lenet()
     model.load_state_dict(torch.load('model.pt'))
     model.eval()
     return model

model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:922abc05756421cbdffc98dc27a3eef6abf476761170e77656e7bb5477a45b10
-size 573263
+oid sha256:72ef238f2653e7e2d135b3ecee92ecc02ba463f99dd1871dd9735692a60b3a10
+size 249671