Thanks for the update to tf2. Tested the code and found the new modifications seems to break. Reverting to the old line fixes it. Refer to PR #1
