Cuda Error with Tensorflow > 2.0

I am getting runtime error where the Tensorflow library is unable to initialize Cuda with versions higher than 2.0. Here are the software configs:

Ubuntu 16.04
Cuda: 10.1
Driver: 418.87.01
Cuddn: 7.5.4
Tensorflow: 2.2 (tried with 2.0)

Link to code: test.py · GitHub
Link to output: output.log · GitHub

I have tried using the memory growth approaches described here - https://github.com/tensorflow/tensorflow/issues/24828
Also the code works fine in Tensorflow Version 1.8. Not sure whats going on.

I’ve only run into this issue when I two models training on the same GPU.

Was the output any differet when you gave the memory growth approaches a try?

devices = tf.config.experimental.list_physical_devices('GPU')
for device in devices
    tf.config.experimental.set_memory_growth(devices, True)

I seem to remember that 1.4 was the last version you can use with 2.0. I could be wrong, however; so check it.