Fix changed underlying API

Reviewer: @andyyoung