Return values are forced to be on CPU

[maria-korosteleva]

Hi!
Thanks for the great tool!

I noticed that the output of kmeans() is forced to be on the CPU regardless of what device was requested. Line 121 here

return choice_cluster.cpu(), initial_state.cpu()

I'd expect the return values to be on the same device as the device requested in the function call.

Is there a particular reason for forcing the cpu in returned values?

[subhadarship]

Thanks for the Q! There is no necessary reason to force the output to be on the CPU. I just did it because it fit my use case. Feel free to change that code to fit your use case (or even contribute by making a pull request!)