Hi all, I am trying to implement a variation 2 of torch.max() and confused about the implementation of dim parameter. Here is an example,
inp = torch.arange(0., 24).view(2, 3, 4)
torch.max(inp, dim=1)

torch.return_types.max(
values=tensor([[ 8.,  9., 10., 11.],
        [20., 21., 22., 23.]]),
indices=tensor([[2, 2, 2, 2],
        [2, 2, 2, 2]]))

So, the max value within each colored rectangle is returned as in the figure.

Likewise, torch.max(inp, dim=2) returns the max value of each colored rectangle in the below figure,
torch.return_types.max(
values=tensor([[ 3.,  7., 11.],
        [15., 19., 23.]]),
indices=tensor([[3, 3, 3],
        [3, 3, 3]]))


My question is how can I reshape the input as done with the torch.max() in order to use dim parameter in my custom max implementation ? I’ve experimented view with different values but it didn’t work so far.