fix test_data_loader on Windows and macOS (#5912)

* fix test_data_loader on Windows and macOS * Update test/test_prototype_builtin_datasets.py Co-authored-by: Nicolas Hug <[email protected]> Co-authored-by: Nicolas Hug <[email protected]>
pytorch · Apr 28, 2022 · 734ee25 · 734ee25
1 parent a945ec6
commit 734ee25
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/test/test_prototype_builtin_datasets.py b/test/test_prototype_builtin_datasets.py
@@ -125,6 +125,12 @@ def test_serializable(self, test_home, dataset_mock, config):
 
         pickle.dumps(dataset)
 
+    # This has to be a proper function, since lambda's or local functions
+    # cannot be pickled, but this is a requirement for the DataLoader with
+    # multiprocessing, i.e. num_workers > 0
+    def _collate_fn(self, batch):
+        return batch
+
     @pytest.mark.parametrize("num_workers", [0, 1])
     @parametrize_dataset_mocks(DATASET_MOCKS)
     def test_data_loader(self, test_home, dataset_mock, config, num_workers):
@@ -135,7 +141,7 @@ def test_data_loader(self, test_home, dataset_mock, config, num_workers):
             dataset,
             batch_size=2,
             num_workers=num_workers,
-            collate_fn=lambda batch: batch,
+            collate_fn=self._collate_fn,
         )
 
         next(iter(dl))