Little fix of code