From 57eb6ecea7695dbc9e5abbbc5695201982b00da5 Mon Sep 17 00:00:00 2001
From: Anthony Wang <ta180m@pm.me>
Date: Wed, 25 Aug 2021 21:01:46 -0500
Subject: [PATCH] Edits to make the script actually compile and achieve 99% on
 MNIST

---
 mnist.py | 23 +++++++++++------------
 1 file changed, 11 insertions(+), 12 deletions(-)

diff --git a/mnist.py b/mnist.py
index 2b4085f..f6aba1d 100644
--- a/mnist.py
+++ b/mnist.py
@@ -6,6 +6,7 @@ from torchvision import datasets
 from torchvision.transforms import ToTensor, Lambda, Compose
 import matplotlib.pyplot as plt
 
+
 training_data = datasets.MNIST(
     root="data",
     train=True,
@@ -20,7 +21,7 @@ test_data = datasets.MNIST(
     transform=ToTensor(),
 )
 
-batch_size = 64
+batch_size = 100
 
 train_loader = DataLoader(training_data, batch_size=batch_size)
 test_loader = DataLoader(test_data, batch_size=batch_size)
@@ -47,16 +48,15 @@ class CNN(nn.Module):
         self.fc2 = nn.Linear(in_features=600, out_features=120)
         self.fc3 = nn.Linear(in_features=120, out_features=10)
 
-
-def forward(self, x):
-    out = self.layer1(x)
-    out = self.layer2(out)
-    out = out.view(out.size(0), -1)
-    out = self.fc1(out)
-    out = self.drop(out)
-    out = self.fc2(out)
-    out = self.fc3(out)
-    return out
+    def forward(self, x):
+        out = self.layer1(x)
+        out = self.layer2(out)
+        out = out.view(out.size(0), -1)
+        out = self.fc1(out)
+        out = self.drop(out)
+        out = self.fc2(out)
+        out = self.fc3(out)
+        return out
 
 
 model = CNN()
@@ -93,7 +93,6 @@ for epoch in range(num_epochs):
             total = 0
             correct = 0
             for images, labels in test_loader:
-                images, labels = images.to(device), labels.to(device)
                 labels_list.append(labels)
             
                 test = Variable(images.view(batch_size, 1, 28, 28))