trained pth files 2

app.py

@@ -8,7 +8,7 @@ import torch.optim as optim
 
 # This is just to show an interface where one draws a number and gets prediction. 
 
-n_epochs = 100
+n_epochs = 10
 batch_size_train = 128
 batch_size_test = 1000
 learning_rate = 0.01
@@ -67,16 +67,16 @@ class MNIST_Model(nn.Module):
         return F.log_softmax(x)
 
 train_loader = torch.utils.data.DataLoader(
-  torchvision.datasets.MNIST('.files/', train=True, download=True,
+  torchvision.datasets.MNIST('files/', train=True, download=True,
                              transform=torchvision.transforms.Compose([
                                torchvision.transforms.ToTensor(),
                                torchvision.transforms.Normalize(
-                                 (0.1307,), (0.3081,))
+                                 mean=(0.1307,), std=(0.3081,))
                              ])),
   batch_size=batch_size_train, shuffle=True)
 
 test_loader = torch.utils.data.DataLoader(
-  torchvision.datasets.MNIST('.files/', train=False, download=True,
+  torchvision.datasets.MNIST('files/', train=False, download=True,
                              transform=torchvision.transforms.Compose([
                                torchvision.transforms.ToTensor(),
                                torchvision.transforms.Normalize(
@@ -84,25 +84,24 @@ test_loader = torch.utils.data.DataLoader(
                              ])),
   batch_size=batch_size_test, shuffle=True)
 
-def train(epochs,network,optimizer,train_loader):
+def train(epoch,network,optimizer,train_loader):
     
     train_losses=[]
     network.train()
-    for epoch in range(epochs):
-        for batch_idx, (data, target) in enumerate(train_loader):
-            optimizer.zero_grad()
-            output = network(data)
-            loss = F.nll_loss(output, target)
-            loss.backward()
-            optimizer.step()
-            if batch_idx % log_interval == 0:
-                print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
-                epoch, batch_idx * len(data), len(train_loader.dataset),
-                100. * batch_idx / len(train_loader), loss.item()))
-                train_losses.append(loss.item())
-        
-                torch.save(network.state_dict(), MODEL_WEIGHTS_PATH)
-                torch.save(optimizer.state_dict(), OPTIMIZER_PATH)
+    for batch_idx, (data, target) in enumerate(train_loader):
+        optimizer.zero_grad()
+        output = network(data)
+        loss = F.nll_loss(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+            epoch, batch_idx * len(data), len(train_loader.dataset),
+            100. * batch_idx / len(train_loader), loss.item()))
+            train_losses.append(loss.item())
+    
+            torch.save(network.state_dict(), MODEL_WEIGHTS_PATH)
+            torch.save(optimizer.state_dict(), OPTIMIZER_PATH)
 
 def test():
     test_losses=[]
@@ -121,6 +120,7 @@ def test():
         acc = acc.item()
         test_metric = '〽Current test metric -> Avg. loss: `{:.4f}`, Accuracy: `{:.0f}%`\n'.format(
         test_loss,acc)
+        print(test_metric)
         return test_metric,acc
 
 
@@ -143,8 +143,11 @@ if os.path.exists(model_state_dict) and os.path.exists(optimizer_state_dict):
     optimizer_state_dict = torch.load(optimizer_state_dict)
     optimizer.load_state_dict(optimizer_state_dict)   
 # Train
-train(n_epochs,network,optimizer,train_loader)
-test()
+
+for epoch in range(n_epochs):
+
+  train(epoch,network,optimizer,train_loader)
+  test()
 
 
 def image_classifier(inp):