edited code to fix bug