Make tests compatible between tf 1 and 2.

5 jobs for dev_tf2 in 39 seconds
Status Job ID Name Coverage
  Lint
passed #302978
docker
lint

00:00:16

 
  Test
passed #302980
docker
unittests_tf1_py3

00:00:22

passed #302982
docker
unittests_tf2_py3

00:00:22

passed #302979
docker
unittest_tf1_py2

00:00:20

passed #302981
docker
unittest_tf2_py2

00:00:22