Revert "Revert "bug fix""

This reverts commit bf94ab06a7ea13e053d08a4420282948c1c6c314.