Fix minor bug returning

Test 1