Made test_ndarray_batch_source python 3 compatible.

mila-iqia · Dec 4, 2015 · d683774 · d683774
1 parent 32a7797
commit d683774
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/transformers/test_image.py b/tests/transformers/test_image.py
@@ -319,13 +319,13 @@ def test_ndarray_batch_source(self):
         for loc in [(0, 0), (0, 1), (1, 0), (1, 1)]:
             stream = FixedSizeCrop(self.batch_stream, (5, 4),
                                    which_sources=('source1',), location=loc)
-            batch = stream.get_epoch_iterator().next()
-            assert batch[0].shape[1:] == (3, 5, 4)
-            assert batch[0].shape[0] in (1, 2)
             # seen indices should only be of that length in after last location
             if 3 * 7 * 5 == len(seen_indices):
                 assert False
-            seen_indices = numpy.union1d(seen_indices, batch[0].flatten())
+            for batch in stream.get_epoch_iterator():
+                assert batch[0].shape[1:] == (3, 5, 4)
+                assert batch[0].shape[0] in (1, 2)
+                seen_indices = numpy.union1d(seen_indices, batch[0].flatten())
         assert 3 * 7 * 5 == len(seen_indices)
 
     def test_list_batch_source(self):