Fix build

I broke it while cherry picking the last commit.