Problem with CuDNN initialization

[jeremy]

Try adding the following after importing Tensorflow:

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