Кастомная функция потерь keras

Есть желание использовать модель другого типа в качестве "последнего слоя" нейросети. Пришла идея зашить модель в функцию потерь, однако, как я поняла, для keras моделей производная от функции берется через метод GradientTape, который очень похож почти на аналитический, и требует именно функций. По идее, градиент можно вычислить и численно, но тогда нужно переписывать с нуля оптимизатор?