Revert "Revert "bug fix""

This reverts commit e4dd83218383beb9b3294d0790a767699cb40284.