Incorrect step assignments with custom step function

This is great - thanks a lot!