Small improve of code