minor code change

added 4 lines of code!