Update tensorflow2 code
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/61406/checks?check_run_id=15384363955) 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. Hi @pkanwar23 Can you please review this PR ? Thank you!
In this version, we use tf.compat.v1.enable_v2_behavior() and tf.compat.v1.disable_v2_behavior() instead of directly using _pywrap_tf2. Additionally, the function is_tf2_behavior_enabled() now uses tf.executing_eagerly() instead of _pywrap_tf2.is_enabled(), which is the recommended way to check if TensorFlow 2.0 behavior is enabled.