Keras. Загрузка датасета в видеопямать

Всем привет. До недавнего времени пользовался старой версией keras без tensorflow backend-а, перешел на актуальную и столкнулся с проблемой. Весь датасет у меня помещается в ram, но в объеме, большем, чем объем видеопамяти на моих gpu. В старых версиях я вызывал

model.fit(train_x, train_y, epochs=120, batch_size=batch_size, validation_data=(val_x, val_y)

где train_* и val_* - это numpy массивы, и все работало как надо. Т.е keras как-то по-умному грузил датасет в видеопамять кусками. Теперь же получаю ошибку,

2022-09-28 10:52:52.463643: W tensorflow/core/common_runtime/bfc_allocator.cc:479] Allocator (GPU_0_bfc) ran out of memory trying to allocate 5.84GiB (rounded to 6276710400)requested by op _EagerConst

. Очевидно, он пытается загрузить весь датасет в память разом. Как избежать этого в актуальной версии без использования Sequence?