Bug fix

Had the bug fix in the wrong place, ended up causing another bug. Moved
it to the appropriate place, no longer an issue