max operation for specific elements in 4 dimensional array in parallel way

Question

I am coding pytorch. Between the torch inference code, I add some peripheral code for my own interest. This code works fine, but it is too slow. The reason might be for iter