diff --git a/py/facerec/feature.py b/py/facerec/feature.py
index de746f5..34c86fb 100644
--- a/py/facerec/feature.py
+++ b/py/facerec/feature.py
@@ -113,7 +113,7 @@ def compute(self, X, y):
         y = np.asarray(y)
         # calculate dimensions
         d = XC.shape[0]
-        c = len(np.unique(y))        
+        c = len(np.unique(y))
         # set a valid number of components
         if self._num_components <= 0:
             self._num_components = c-1
@@ -124,7 +124,7 @@ def compute(self, X, y):
         # calculate the within and between scatter matrices
         Sw = np.zeros((d, d), dtype=np.float32)
         Sb = np.zeros((d, d), dtype=np.float32)
-        for i in range(0,c):
+        for i in np.unique(y):
             Xi = XC[:,np.where(y==i)[0]]
             meanClass = np.mean(Xi, axis = 1).reshape(-1,1)
             Sw = Sw + np.dot((Xi-meanClass), (Xi-meanClass).T)