Minor bug fix prior to refraction of code