Fixed problems in code