Переобучение нейронной сети VGG16 keras

Мне нужно переобучить сеть VGG16 для распознавания моего собственного класса изображений, при загрузке данных происходит ошибка. С чем это может быть связано?

from keras.applications import vgg16
import os
import numpy as np
from keras_preprocessing.image import ImageDataGenerator

conv_base = vgg16.VGG16(weights='imagenet', include_top=False)

base_dir = "train_dod"
train_dir = os.path.join(base_dir, "train")
validation_dir = os.path.join(base_dir, "validation")
test_dir = os.path.join(base_dir, "test")

datagen = ImageDataGenerator(rescale=1./255)
batch_size = 20

def extract_features(directory, sample_count):
    features = np.zeros(shape=(sample_count, 4, 4, 512))
    labels = np.zeros(shape=(sample_count))
    generator = datagen.flow_from_directory(
        directory,
        target_size=(150, 150),
        batch_size=batch_size,
        class_mode='binary'
    )
    i = 0
    for inputs_batch, labels_batch in generator:
        features_batch = conv_base.predict(inputs_batch)
        features[i * batch_size : (i + 1) * batch_size] = features_batch
        labels[i * batch_size : (i + 1) * batch_size] = labels_batch
        i += 1
        if i * batch_size >= sample_count:
            break
    return features, labels

train_features, train_labels = extract_features(train_dir, 900)
validation_features, validation_labels = extract_features(validation_dir, 100)
test_features, test_labels = extract_features(test_dir, 779)

train_features = np.reshape(train_features, (900, 4 * 4 * 512))
validation_features = np.reshape(validation_features, (100, 4 * 4 * 512))
test_features = np.reshape(test_features, (779, 4 * 4 * 512))

Получаем ошибку:

Traceback (most recent call last):
  File "D:\Study\Progs\testik\venv\lib\site-packages\keras\engine\data_adapter.py", line 1198, in _truncate_execution_to_epoch
    self._steps_per_execution.assign(self._inferred_steps)
AttributeError: 'int' object has no attribute 'assign'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "D:\Study\Progs\testik\testnet.py", line 35, in <module>
    train_features, train_labels = extract_features(train_dir, 900)
  File "D:\Study\Progs\testik\testnet.py", line 27, in extract_features
    features_batch = conv_base.predict(inputs_batch)
  File "D:\Study\Progs\testik\venv\lib\site-packages\keras\engine\training.py", line 1747, in predict
    for _, iterator in data_handler.enumerate_epochs():  # Single epoch.
  File "D:\Study\Progs\testik\venv\lib\site-packages\keras\engine\data_adapter.py", line 1179, in enumerate_epochs
    with self._truncate_execution_to_epoch():
  File "C:\Users\Ilya\AppData\Local\Programs\Python\Python39\lib\contextlib.py", line 117, in __enter__
    return next(self.gen)
  File "D:\Study\Progs\testik\venv\lib\site-packages\keras\engine\data_adapter.py", line 1203, in _truncate_execution_to_epoch
    self._steps_per_execution.assign(original_value)
AttributeError: 'int' object has no attribute 'assign'

Process finished with exit code 1