The custom max should return the indices of all maximum values instead of the first one being encountered as in torch.max. I want to add dim as a parameter to my custom max, like in torch.max. But I don’t know how to implement the dimension reduction effect controlled by dim.
Documentation 6 of torch.max states as below but the input is not in the shape of (Ax1xB), so using torch.squeeze does not yield to reduced dimensions.

Otherwise, dim is squeezed (see torch.squeeze()), resulting in the output tensors having 1 fewer dimension than input