Update return type