Minor bug fix for return