Update model.py

model.py

@@ -110,6 +110,10 @@ def download_data():
     # 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)
+
 
 # uses GPU if available
 if torch.cuda.is_available():
@@ -119,9 +123,6 @@ else:
 
 device = torch.device(dev)
 
-# gets mean and std of dataset
-mean, std = get_mean_std(train_loader)
-
 
 def run_model():
     # defines parameters