Fix return register handling