def convert_data(data):
    """ Helper function to convert data to the right shape. """
    shape = data.shape
    d = len(shape)
    if d == 2:
        m = shape[0]
        n = shape[1]
        data = data.reshape(1, m, n)
    elif d == 3:
        m = shape[1]
        n = shape[2]
    else:
        raise ValueError(
            'Could not bring data with shape {} into right                          form.'
            .format(shape))
    return data, d, m, n
