Fixed problems with code