pytorch tmp (aws#382)

Author: NihalHarish

tests/zero_code_change/pt_utils.py

@@ -7,22 +7,34 @@
 import torch.nn.functional as F
 import torchvision
 import torchvision.transforms as transforms
+from packaging import version
 
 
 def get_dataloaders() -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
     transform = transforms.Compose(
         [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
     )
 
+    # Temporary Change to allow the test to run with pytorch 1.7 RC3
+    # Smdebug breaks when num_workers>0 for Pytorch 1.7.0
+    if version.parse(torch.__version__) >= version.parse("1.7.0"):
+        num_workers = 0
+    else:
+        num_workers = 2
+
     trainset = torchvision.datasets.CIFAR10(
         root="./data", train=True, download=True, transform=transform
     )
-    trainloader = torch.utils.data.DataLoader(trainset, batch_size=4, shuffle=True, num_workers=2)
+    trainloader = torch.utils.data.DataLoader(
+        trainset, batch_size=4, shuffle=True, num_workers=num_workers
+    )
 
     testset = torchvision.datasets.CIFAR10(
         root="./data", train=False, download=True, transform=transform
     )
-    testloader = torch.utils.data.DataLoader(testset, batch_size=4, shuffle=False, num_workers=2)
+    testloader = torch.utils.data.DataLoader(
+        testset, batch_size=4, shuffle=False, num_workers=num_workers
+    )
 
     classes = ("plane", "car", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck")
     return trainloader, testloader