Update model.py

model.py

@@ -101,19 +101,14 @@ class Net(nn.Module):
         return x
 
 
-def download_data():
-    # downloads and loads MNIST train set
-    transform = transforms.Compose([transforms.ToTensor(), transforms.RandomAffine(degrees=10, translate=(0.1,0.1))])
-    train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
-    train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True, pin_memory=True)
-    
-    # downloads and loads MNIST test set
-    val_data = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
-    val_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=False, pin_memory=True)
-    
-    # gets mean and std of dataset
-    mean, std = get_mean_std(train_loader)
+# downloads and loads MNIST train set
+transform = transforms.Compose([transforms.ToTensor(), transforms.RandomAffine(degrees=10, translate=(0.1,0.1))])
+train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
+train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True, pin_memory=True)
 
+# downloads and loads MNIST test set
+val_data = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
+val_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=False, pin_memory=True)
 
 # uses GPU if available
 if torch.cuda.is_available():
@@ -123,6 +118,9 @@ else:
 
 device = torch.device(dev)
 
+# gets mean and std of dataset
+mean, std = get_mean_std(train_loader)
+
 
 def run_model():
     # defines parameters