Fix some issues with the previous change.