Fix error with argument handling