fix bug in train_NB method