Fix flow matching function