Code modification

Bug fix