Moved the "Fix" of wrong delegate type

The While loop code was called from other places, so moved the addition
of the Unit type up to there.