ValueError: обнаружены входные переменные с несовместимым количеством выборок: [84, 5]

Question

Я делаю модель CNN, используя VGG16 и несколько настраиваемых слоев в конце. Набор данных выглядит следующим образом:

Found 200 images belonging to 2 classes. 
Found 84 images belonging to 2 classes.

гиперпараметры:

epochs= 25
lr = 1e-4
BS = 16

Модель выглядит следующим образом:

basemodel = VGG16(weights="imagenet", include_top=False,input_tensor=Input(shape=(224, 224, 3)))
headModel = basemodel.output
headModel = AveragePooling2D(pool_size=(4, 4))(headModel)#pool_size=(4, 4)
headModel = Flatten(name="flatten")(headModel)
#headModel = Dropout(0.5)(headModel)
headModel = Dense(256, activation="relu")(headModel)
headModel = Dropout(0.5)(headModel)
#headModel = Dense(256, activation="relu")(headModel)
#headModel = Dropout(0.3)(headModel)
headModel = Dense(2, activation="softmax")(headModel)




from sklearn.metrics import confusion_matrix
Y_pred = model.predict_generator(test_generator, np.ceil(len(test_generator) // BS))
y_pred = np.argmax(Y_pred, axis=1)
print(confusion_matrix(test_generator.classes, y_pred))

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-65-29f820d7a706> in <module>
----> 1 print(confusion_matrix(test_generator.classes, y_pred))

~/opt/anaconda3/lib/python3.7/site-packages/sklearn/metrics/_classification.py in confusion_matrix(y_true, y_pred, labels, sample_weight, normalize)
    266 
    267     """
--> 268     y_type, y_true, y_pred = _check_targets(y_true, y_pred)
    269     if y_type not in ("binary", "multiclass"):
    270         raise ValueError("%s is not supported" % y_type)

~/opt/anaconda3/lib/python3.7/site-packages/sklearn/metrics/_classification.py in _check_targets(y_true, y_pred)
     78     y_pred : array or indicator matrix
     79     """
---> 80     check_consistent_length(y_true, y_pred)
     81     type_true = type_of_target(y_true)
     82     type_pred = type_of_target(y_pred)

~/opt/anaconda3/lib/python3.7/site-packages/sklearn/utils/validation.py in check_consistent_length(*arrays)
    210     if len(uniques) > 1:
    211         raise ValueError("Found input variables with inconsistent numbers of"
--> 212                          " samples: %r" % [int(l) for l in lengths])
    213 
    214 

ValueError: Found input variables with inconsistent numbers of samples: [84, 5]

Я получаю эту ошибку в большинстве случаев оценочных функций. Пожалуйста, дайте мне знать, что я делаю не так и требуется ли что-то еще из кода.