Update missing dependency

Author: borchero

pycave/bayes/gmm.py

@@ -3,11 +3,11 @@
 import torch
 import torch.nn as nn
 import torch.distributions as dist
+from sklearn.cluster import KMeans
 import pyblaze.nn as xnn
 from pyblaze.utils.stdio import ProgressBar
-from pyblaze.utils.torch import to_one_hot
-from sklearn.cluster import KMeans
-from .utils import log_normal, log_responsibilities, max_likeli_means, max_likeli_covars
+from .utils import log_normal, log_responsibilities, max_likeli_means, max_likeli_covars, \
+    to_one_hot
 
 class GMMConfig(xnn.Config):
     """

pycave/bayes/utils.py

@@ -193,3 +193,22 @@ def power_iteration(A, eps=1e-7, max_iterations=100):
             break
 
     return v
+
+
+def to_one_hot(X, n):
+    """
+    Creates a one-hot matrix from a set of indices.
+
+    Parameters:
+    -----------
+    - X: torch.Tensor [N, D]
+        The indices to convert into one-hot vectors.
+    - n: int
+        The number of entries in the one-hot vectors.
+
+    Returns:
+    --------
+    - torch.Tensor [N, D, n]
+        The one-hot matrix.
+    """
+    return torch.eye(n, device=X.device)[X]