After digging through the code I find that the board type is set ti ICE_V1. If you add this function:
UTILsSetBoardType(AM335X_BOARD_TYPE_ICE_V2);
It will set it to the correct board type.
I also had to add:
OLEDDelay(OledDelay);
between the OLED commands.
After I added this it worked fine.