K-Means Algorithm using Pytorch and CUDA