Revert "Address PR comments."

This reverts commit 168289112a520a0dce836df96f8a908772036b73.
Address PR comments.