Add missing arg type on memoize_on_first_arg