Fix multiprocessing in torchvision_tutorial.py

Author: alperenunlu

Diff for: intermediate_source/torchvision_tutorial.py

@@ -381,15 +381,19 @@ def get_transform(train):
 # Before iterating over the dataset, it's good to see what the model
 # expects during training and inference time on sample data.
 import utils
+import platform
 
+# In Windows and MacOS, the number of workers have to be set to 0 if we are using notebook
+# We can set num_workers > 0 if we are using a python script and wrapping the code in if __name__ == "__main__": block
+num_workers = 0 if platform.system() in ('Windows', 'Darwin') else 4
 
 model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights="DEFAULT")
 dataset = PennFudanDataset('data/PennFudanPed', get_transform(train=True))
 data_loader = torch.utils.data.DataLoader(
     dataset,
     batch_size=2,
     shuffle=True,
-    num_workers=4,
+    num_workers=num_workers,
     collate_fn=utils.collate_fn
 )
 
@@ -428,20 +432,22 @@ def get_transform(train):
 dataset = torch.utils.data.Subset(dataset, indices[:-50])
 dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
 
+num_workers = 0 if platform.system() in ('Windows', 'Darwin') else 4
+
 # define training and validation data loaders
 data_loader = torch.utils.data.DataLoader(
     dataset,
     batch_size=2,
     shuffle=True,
-    num_workers=4,
+    num_workers=num_workers,
     collate_fn=utils.collate_fn
 )
 
 data_loader_test = torch.utils.data.DataLoader(
     dataset_test,
     batch_size=1,
     shuffle=False,
-    num_workers=4,
+    num_workers=num_workers,
     collate_fn=utils.collate_fn
 )