Minor update to previous update

-Made code look a tad nicer