Minor change to code structure