Code modification

bug fix