Code clean up, minor fix.

Minor formatting and fixed import