Optical Flow SDK (Turing/Ampere) + OpenCV Efficiently Convert Vectors to Magnitude

Hiya, we got started with using the OpenCV integration of the new V2.0 Optical Flow SDK in Python and it looks great. Was wondering if Nvidia have any advice or sample implementations of an efficient way to convert it to a black-and-white magnitude field in Python?

A simple numpy operation does the job (flow_magnitude = np.linalg.norm(flow, axis=2)) but becomes a performance bottleneck.

Thanks for any help,