Как исправить перегруз GPU-памяти с Tensorflow?

Написал модель на Keras, уменьшил изображения, но запускаю код и еще до обучения модели все виснет и выключается комп. Я думаю это из-за GPU-памяти (перегруз), но как исправить не знаю (уменьшил батч - не помогло). Но если я отключу GPU, то на СPU все работает (только медленно). Помогите пж!