[fix] fix bug of argmax.py when the dim-th dimension is less than BLOCK_M

[ConvolutedDog]

PR Category

Operator

Type of Change

Bug Fix

Description

Fix bug of argmax.py when the dim-th dimension is less than BLOCK_M.

Issue

None

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

None

[StrongSpoon]

function argmax works well on my device. could you please show me what the bug is?

[ConvolutedDog]

Can you try to get into this example, I'm not sure if it's a bug with my env:

[StrongSpoon]

It's a bug and I have reproduced it. This is caused by the improper parallelism. When the input tensor is of shape [3, 64, 55, 55], the variable K is 193600, and it exceeds the upper bound of grid size on the second dimension.
In my view, the resolution is to design a better lambda function of grid, thus allowing larger size of K.