Improving compatibility between TF1 and TF2

[Enzodtz]

Description

Some imports were changed in TF1 to TF2. I was using TF2 with this module and got problems, so I fixed it. Just 5 lines of code.
I don't know if this is the best way to check the version.

Type of change

For a new feature or function, please create an issue first to discuss it
with us before submitting a pull request.

Note: Please delete options that are not relevant.

• Bug fix (non-breaking change which fixes an issue)
• Documentation update
• [ x ] TensorFlow 2 migration
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• A new research paper code implementation
• Other (Specify)

Tests

None, it's a pretty simple change.

Checklist

• [ x ] I have signed the Contributor License Agreement.
• [ x ] I have read guidelines for pull request.
• [ x ] My code follows the coding guidelines.
• [ x ] I have performed a self code review of my own code.
• I have commented my code, particularly in hard-to-understand areas.
• I have made corresponding changes to the documentation.
• [ x ] My changes generate no new warnings.
• I have added tests that prove my fix is effective or that my feature works.