Fixing issues introduced by changes