Small adjustment of code