Cleaned up the edit

Moved a little bit of code around