Re-added forgotten code
Before my last commit I forgot to change some things back after testing,
this is now fixed and it's all working as intended again.