Update type from #1124 in PAI
Merge pull request #1125 from cosmoharrigan/pai

Update type from #1124 in PAI