Sandbox run src/cnn.py

sweepai · Nov 26, 2023 · 6f6638f · 6f6638f
1 parent d5b3bff
commit 6f6638f
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/cnn.py b/src/cnn.py
@@ -30,6 +30,7 @@ def forward(self, x):
         output = nn.functional.log_softmax(x, dim=1)
         return output
 
+
 def train(network, train_loader, optimizer):
     network.train()
     for batch_idx, (data, target) in enumerate(train_loader):
@@ -39,13 +40,13 @@ def train(network, train_loader, optimizer):
         loss.backward()
         optimizer.step()
 
+
 if __name__ == "__main__":
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize((0.5,), (0.5,))
-    ])
+    transform = transforms.Compose(
+        [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
+    )
 
-    trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
+    trainset = datasets.MNIST(".", download=True, train=True, transform=transform)
     trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
 
     network = CNN()