Fix tf2.4 (#9120)
* Fix tests for TF 2.4
* Remove <2.4 limitation
* Add version condition
* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
* Update tests/test_optimization_tf.py
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>