Updated code

Bug fix