fix(container): NdarrayContainer batch_to_singles fix
@@ -91,7 +91,8 @@ class NdarrayContainer(DataContainer["np.ndarray", "np.ndarray"]):
def batch_to_singles(cls, batch, batch_axis=0):
import numpy as np
- return np.split(batch, batch.shape[batch_axis], axis=batch_axis)
+ return [np.squeeze(arr, axis=batch_axis) for arr in
+ np.split(batch, batch.shape[batch_axis], axis=batch_axis)]
@classmethod
@inject
