Fix AttriburteError when call model.export
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/tensorflow/tensorflow/pull/62733/checks?check_run_id=20149111647) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request. Check out this pull request on&nbsp; <a href="https://app.reviewnb.com/tensorflow/tensorflow/pull/62733"><img align="absmiddle"  alt="ReviewNB" height="28" class="BotMessageButtonImage" src="https://raw.githubusercontent.com/ReviewNB/support/master/images/button_reviewnb.png"/></a> 

 See visual diffs & provide feedback on Jupyter Notebooks. 

---

 <i>Powered by <a href='https://www.reviewnb.com/?utm_source=gh'>ReviewNB</a></i> Hi @ziyeqinghan Can you please review this PR ? Thank you! Hi @ziyeqinghan Can you please review this PR ? Thank you! LGTMed the PR.
![2024-01-04-115332_1059x87_scrot](https://github.com/tensorflow/tensorflow/assets/68238841/9a23e905-c1c6-4846-a180-be4d967a9cb6)
