Only stucks at the end and I think I know the fix