Изменил подход к управлению памятью, что снизило потребление в 2 раза. А это в свою очередь позволило тренировать на большем разрешении и увеличить размер пакета. На изображениях результат 24 часов тренировки. Утечек памяти не было.