Как я могу распараллелить map_fn Tensorflow в GPU?

не могу опубликовать рассматриваемый код, но я постараюсь описать свой сценарий.

У меня есть настраиваемая операция тензорного потока, которую я определил с помощью автографа Tensorflow. Он содержит множество операций управления тензорным потоком, автоматически конвертируемых с помощью автографа. Я запустил эту операцию на большом количестве пакетов и понял, что она работает очень медленно на графическом процессоре, потому что она не распараллеливается через map_fn, но на CPU она выполняется очень быстро и распараллеливается. Я подтвердил, что процессор распараллеливается, а графический процессор не прикрепляет tf.Print и начало и конец операции, если одновременно выполняется только одна операция, «Начало» и «Конец» будут отображаться по порядку, в противном случае они бы не стали.

В общем, есть ли идеи, почему моя операция будет автоматически распараллеливаться на CPU, но не на GPU, учитывая тот факт, что у него много операций потока управления? Я также уже пробовал вручную установить параметр "parallel_iterations" для map_fn.

Спасибо.