squeeze_batch_dim test sketch

Author: cmdupuis3

xbatcher/tests/test_generators.py

@@ -160,3 +160,23 @@ def test_batch_3d_2d_input_concat(sample_ds_3d, bsize):
             * (sample_ds_3d.dims['y'] // bsize)
             * sample_ds_3d.dims['time']
         )
+
+
+@pytest.mark.parametrize('bsize', [5, 10])
+def test_batch_3d_squeeze_batch_dim(sample_ds_3d, bsize):
+    xbsize = 20
+    bg = BatchGenerator(
+        sample_ds_3d,
+        input_dims={'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},
+        squeeze_batch_dim=True,
+    )
+    for ds_batch in bg:
+        assert ds_batch['x'].shape == [bsize, xbsize]