Sandbox run src/cnn.py

sweepai · Nov 27, 2023 · 53a51f8 · 53a51f8
1 parent 7cfe902
commit 53a51f8
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/cnn.py b/src/cnn.py
@@ -36,12 +36,12 @@ def train_cnn(self, trainloader, epochs=3):
 
         torch.save(self.state_dict(), "mnist_cnn_model.pth")
 
-transform = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
-])
 
-trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
+transform = transforms.Compose(
+    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
+)
+
+trainset = datasets.MNIST(".", download=True, train=True, transform=transform)
 trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
 
 cnn = CNN()