Fix topk IMA (#93095)
Hopefully, this will fix https://github.com/pytorch/pytorch/issues/93006. ~I can not reproduce that issue: I can catch the IMA with compute sanitizer on nightly build, but not on source build of master. So there is no way for me to validate if my fix is correct or not.~ Edit: Thanks for the help of @ptrblck, this fix is validated.
But by reading the code, I believe this is a similar issue as https://github.com/pytorch/pytorch/pull/83042, so I apply the same fix for `mbtopk::gatherTopK`. We can wait until tomorrow's nightly build to see if #93006 disappear.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/93095
Approved by: https://github.com/ngimel
