Some code cleanup in an example

Some code cleanup in an example