Find index of max value of a 3-dimensional array (for each pixel)

[miguelgrine]

I'm translating MATLAB code into C++ using OpenCV. I have an image for which I computed the gradient and then the magnitude like:

  if(size(I,3) == 3)
    [DX,DY,~] = gradient(I);
    mag = DX .^ 2 + DY .^ 2;
    [~,channel] = max(mag,[],3);
    dx = DX(:,:, 1) .* (channel == 1) + DX(:,:, 2) .* (channel == 2) + DX(:,:, 3) .* (channel == 3);
    dy = DY(:,:, 1) .* (channel == 1) + DY(:,:, 2) .* (channel == 2) + DY(:,:, 3) .* (channel == 3);
    complex_g = complex(dx,dy);
else
    [dx,dy] = gradient(I);
    complex_g = complex(dx,dy);
end

My problem is how to get the max value of the matrix mag for each pixel as I'm doing with [~,channel] = max(mag,[],3); . Can you help me?