now balanced batch size iterator also returning batches with 4 dimens…

…ions
sunyucheng · Jul 4, 2017 · 7990df5 · 7990df5
1 parent e61da63
commit 7990df5
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/braindecode/datautil/iterators.py b/braindecode/datautil/iterators.py
@@ -77,7 +77,13 @@ def get_batches(self, dataset, shuffle):
                                        rng=self.rng,
                                        shuffle=shuffle)
         for batch_inds in batches:
-            yield (dataset.X[batch_inds], dataset.y[batch_inds])
+            batch_X = dataset.X[batch_inds]
+            batch_y = dataset.y[batch_inds]
+
+            # add empty fourth dimension if necessary
+            if batch_X.ndim == 3:
+                batch_X = batch_X[:, :, :, None]
+            yield (batch_X, batch_y)
 
     def reset_rng(self):
         self.rng = RandomState(328774)