Мне реально помогло опускание ниже питонячьих библиотек единожды в жизни, когда пришлось преобразование радона под gpu перефигачивать с возможностью пропускать через него градиент. Там вышло примерно +250% скорости на обучении и +50% на инференсе. Обычно подобная игра свеч не стоит.