Cleaning up some code.

Cleaning up the code