fix for kmeans state updater - NPE. set class index on the dataset