Keras multi GPU in vast.ai

I am trying to run a keras model on vast.ai using multiple GPUs. For that I am using keras.utils.multi_gpu_model , however I keep having this error:

if multi_GPU and n_GPUs  1:
    model = multi_gpu_model(model)

AttributeError: module 'tensorflow_core._api.v2.config' has no attribute 'experimental_list_devices')

I am using this default docker :

Official docker images for deep learning framework TensorFlow

Successfully loaded tensorflow/tensorflow:nightly-gpu-py3

I have also checked the available GPUs and all the GPUs are detected correctly:

Any ideas?

cheers