Small code changes