Fix bugs introduced with changes