add new code

add new code to fix a bug.