Changed the way we count the number of states. I use ass.max() + 1 instead of counting unique entries. This is because in order to count unique entries, I need to also map those unique ID's to contiguous indices