Hi all,
I want to set the k-th smallest value to zero, I get the solution from the network, and it success on my host device(win10, cpu mode), but failed on my server device(ubuntu, gpu mode). I don’t know what’s wrong with my code, could anyone help me?
code:

image.png751×867 44.7 KB

raw tensor:

image.png779×233 9.67 KB

suceess on my win10 and run with cpu mode:

image.png837×521 21.2 KB

failed on my server device and run with gpu mode:

image.png720×272 2.95 KB

I found torch.topk can return the correct indices, but after scatter, the result is wrong on the ubuntu device with gpu mode.