Могу ли я оптимизировать вызов torch::where для быстрого выполнения с помощью других функций? ⇐ C++

[Anonymous]

Сейчас я работаю над проектом, который требует быстрого выполнения и использует много libtorch для необходимых тензорных операций. Во время отладки я заметил, что следующая строка работает очень медленно:

Код: Выделить всё

auto res = torch::where(tensor_A != 0.0, tensor_A - tensor_B, tensor_A);

с tensor_A.sizes() == [1, N, 10, 2], tensor_B.sizes() == [N, 1, 1, 2]< /код>. Очень медленно означает ~15 мс для N=128, что будет максимальным размером N. Как я могу это оптимизировать? Есть ли более быстрая альтернатива с использованием других функций в этом случае приложения?
Заранее спасибо
Я уже пробовал предварительно вычислить условие и вычитание, но получилось не улучшает общую производительность.
Изменить:

Это проект ros2, поэтому я компилирую с использованием сборки colcon и уже включил -03 и -march=native. Я измеряю время с помощью std::chrono::high_solve_lock::now()

Подробнее здесь: https://stackoverflow.com/questions/793 ... -functions