Fixing bug in JS