Fix integration tests

microsoft · Jan 1, 2022 · 6daa4d6 · 6daa4d6
1 parent 7f16553
commit 6daa4d6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/tests/samplers/test_batch.py b/tests/samplers/test_batch.py
@@ -102,6 +102,8 @@ def test_roi(self) -> None:
     @pytest.mark.parametrize("num_workers", [0, 1, 2])
     def test_dataloader(self, sampler: RandomBatchGeoSampler, num_workers: int) -> None:
         ds = CustomGeoDataset()
-        dl = DataLoader(ds, batch_sampler=sampler, num_workers=num_workers)
+        dl = DataLoader(
+            ds, batch_sampler=sampler, num_workers=num_workers, collate_fn=stack_samples
+        )
         for _ in dl:
             continue
diff --git a/tests/samplers/test_single.py b/tests/samplers/test_single.py
@@ -99,7 +99,9 @@ def test_roi(self) -> None:
     @pytest.mark.parametrize("num_workers", [0, 1, 2])
     def test_dataloader(self, sampler: RandomGeoSampler, num_workers: int) -> None:
         ds = CustomGeoDataset()
-        dl = DataLoader(ds, sampler=sampler, num_workers=num_workers)
+        dl = DataLoader(
+            ds, sampler=sampler, num_workers=num_workers, collate_fn=stack_samples
+        )
         for _ in dl:
             continue
 
@@ -147,6 +149,8 @@ def test_roi(self) -> None:
     @pytest.mark.parametrize("num_workers", [0, 1, 2])
     def test_dataloader(self, sampler: GridGeoSampler, num_workers: int) -> None:
         ds = CustomGeoDataset()
-        dl = DataLoader(ds, sampler=sampler, num_workers=num_workers)
+        dl = DataLoader(
+            ds, sampler=sampler, num_workers=num_workers, collate_fn=stack_samples
+        )
         for _ in dl:
             continue