Fixing state issues