[fbsync] fix test_data_loader on Windows and macOS (#5912)

Summary:
* fix test_data_loader on Windows and macOS

* Update test/test_prototype_builtin_datasets.py

Reviewed By: jdsgomes, NicolasHug

Differential Revision: D36095692

fbshipit-source-id: 8835f9a92be235ceec929f2b94ec3edb1946afd1

Co-authored-by: Nicolas Hug <[email protected]>
Co-authored-by: Nicolas Hug <[email protected]>

Author: 2 people

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))