Fix small bug with count matrix construction