Fixing logic issue