Some code tweak