Small code adjustments