Перестроить в Google Colab модель ИИ для работы с большим объемом данных

В Google Colab имеется модель ИИ на BERT. На маленькой выборке она достаточно хорошо считается. Нужная точность достигается.
Но при увеличении объема обучающей выборки начинает не хватать RAM. Нужно перестроить код так, чтобы можно было обучать модель, подгружая большие объемы данных постепенно.