squeeze_batch_dim test (attempt 2)

xarray-contrib · Nov 19, 2021 · 749ac26 · 749ac26
1 parent 0e8f716
commit 749ac26
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/xbatcher/tests/test_generators.py b/xbatcher/tests/test_generators.py
@@ -167,15 +167,15 @@ def test_batch_3d_squeeze_batch_dim(sample_ds_3d, bsize):
     xbsize = 20
     bg = BatchGenerator(
         sample_ds_3d,
-        input_dims={'y': bsize, 'x': xbsize},
+        input_dims={'time': 1, 'y': bsize, 'x': xbsize},
         squeeze_batch_dim=False,
     )
     for ds_batch in bg:
         assert ds_batch['x'].shape == [1, bsize, xbsize]
-        
+
     bg2 = BatchGenerator(
         sample_ds_3d,
-        input_dims={'y': bsize, 'x': xbsize},
+        input_dims={'time': 1, 'y': bsize, 'x': xbsize},
         squeeze_batch_dim=True,
     )
     for ds_batch in bg: