I think auto outp = out.data<float>() and auto idxp = idx.data<int>() should do the trick.
You code is kernel implicitly assuming that you have a contiguous tensor. If you wanted to change that, you could pass PackedAccessors to the kernel. They can be roughly used as arrays. ATen’s native batch norm is an example that uses PackedAccessors 17.
Best regards
Thomas