fixed bug in compute_counts in train_test.R