Update type from #1124 in PAI