[#46][#19][#6]Small Correction

Tested this and it didn't work so I've reverted to the old code which
does work.