Bug fix for onehot encoding