Small code corrections