comments

Author: NihalHarish

tests/zero_code_change/test_pytorch_multiprocessing.py

@@ -41,11 +41,14 @@ def forward(self, x):
 
 def train(rank, model, device, dataloader_kwargs):
     # Training Settings
-    shutil.rmtree(data_dir, ignore_errors=True)
+
     batch_size = 64
     epochs = 1
     lr = 0.01
     momentum = 0.5
+
+    # clear data_dir before saving to it
+    shutil.rmtree(data_dir, ignore_errors=True)
     torch.manual_seed(1 + rank)
     train_loader = torch.utils.data.DataLoader(
         datasets.MNIST(