Yeah, setting a flag in the menu event for selecting the starter would be the more robust option, since it will also work even if the player would release or trade away their starter, for instance - and you'd also not have to deal with extra cases for the starter's evolutions!