Did some minor changes in code