Fixed problems with previous code