Small modification made