Code cleanup

Cleaned up code from previous commit.