def _remove_batch_axis(array):
  """Removes the batch dimension (axis 0) from the array.

  This function is preferred over np.squeeze(array, axis=0) because the
  later only works as expected with numpy arrays. In case the array is a list
  of numpy array, np.squeeze removes all the dimensions with size 1 instead
  of just the 0th dimension. To avoid this undesirable effect, we create a
  separate function.

  Args:
    array: A list of nd numpy arrays or an nd numpy array.

  Returns:
    numpy array with one dimension less than the input.

  Raises:
    ValueError: Raises an error if the array / list doesn't have
        a length 1 along the first (batch) dimension.
  """
  if len(array) != 1:
    raise ValueError('The array doesn\'t have the batch dimension as 1. '
                     'Received an array with length along the batch '
                     'dimension: %d' % len(array))
  return array[0]