Fixing bug with setting state