Update model.py

model.py

@@ -101,14 +101,15 @@ class Net(nn.Module):
         return x
 
 
-# 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)
+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)
 
 # uses GPU if available
 if torch.cuda.is_available():