fix failing tests bcoz of generators

Lightning-AI · Jul 8, 2024 · 575615c · 575615c
1 parent 30a6ce8
commit 575615c
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/streaming/test_dataset.py b/tests/streaming/test_dataset.py
@@ -1046,7 +1046,7 @@ def test_dataset_with_mosaic_mds_data(tmpdir):
     for batch in dataloader:
         assert len(batch["class"]) == 4
         assert len(batch["image"]) == 4
-        assert (_class for _class in batch["class"]) == [4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3]
+        assert list(batch["class"]) == [4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3]
         i += 1
 
     dataloader = DataLoader(dataset, batch_size=4, drop_last=False)
@@ -1056,9 +1056,9 @@ def test_dataset_with_mosaic_mds_data(tmpdir):
             # last batch is smaller than batch_size
             assert len(batch["class"]) == 2
             assert len(batch["image"]) == 2
-            assert (_class for _class in batch["class"]) == [4 * i, 4 * i + 1]
+            assert list(batch["class"]) == [4 * i, 4 * i + 1]
             break
         assert len(batch["class"]) == 4
         assert len(batch["image"]) == 4
-        assert (_class for _class in batch["class"]) == [4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3]
+        assert list(batch["class"]) == [4 * i, 4 * i + 1, 4 * i + 2, 4 * i + 3]
         i += 1