minor bug fix in code