Gaussian Mixture Model Pytorch