Made some changes

Should be ok