Cleaned up the code a little bit.

Signed-off-by: Arica Liu <aricaliu@wpi.edu>