Fix OSCD tests

microsoft · Jan 22, 2023 · 98c690d · 98c690d
1 parent b1a9ae2
commit 98c690d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/datamodules/test_oscd.py b/tests/datamodules/test_oscd.py
@@ -34,7 +34,7 @@ def test_train_dataloader(self, datamodule: OSCDDataModule) -> None:
         batch = next(iter(datamodule.train_dataloader()))
         batch = datamodule.on_after_batch_transfer(batch, 0)
         assert batch["image"].shape[-2:] == batch["mask"].shape[-2:] == (2, 2)
-        assert batch["image"].shape[0] == batch["mask"].shape[0] == 2
+        assert batch["image"].shape[0] == batch["mask"].shape[0] == 1
         if datamodule.bands == "all":
             assert batch["image"].shape[1] == 26
         else:
@@ -47,7 +47,7 @@ def test_val_dataloader(self, datamodule: OSCDDataModule) -> None:
         batch = datamodule.on_after_batch_transfer(batch, 0)
         if datamodule.val_split_pct > 0.0:
             assert batch["image"].shape[-2:] == batch["mask"].shape[-2:] == (2, 2)
-            assert batch["image"].shape[0] == batch["mask"].shape[0] == 2
+            assert batch["image"].shape[0] == batch["mask"].shape[0] == 1
             if datamodule.bands == "all":
                 assert batch["image"].shape[1] == 26
             else:
@@ -59,7 +59,7 @@ def test_test_dataloader(self, datamodule: OSCDDataModule) -> None:
         batch = next(iter(datamodule.test_dataloader()))
         batch = datamodule.on_after_batch_transfer(batch, 0)
         assert batch["image"].shape[-2:] == batch["mask"].shape[-2:] == (2, 2)
-        assert batch["image"].shape[0] == batch["mask"].shape[0] == 2
+        assert batch["image"].shape[0] == batch["mask"].shape[0] == 1
         if datamodule.bands == "all":
             assert batch["image"].shape[1] == 26
         else: