How do I assign the right gpu to tf.keras using ray with allow growth set to true?

def create_keras_model(device_id):
    import tensorflow as tf
    from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
    from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2
    from tensorflow.keras.preprocessing import image
    physical_devices = tf.config.experimental.get_visible_devices(device_type='GPU')
    print(physical_devices[device_id])
    tf.config.experimental.set_visible_devices(physical_devices[device_id], 'GPU')
#     tf.config.experimental.set_memory_growth(physical_devices[device_id], True)
#     tensorflow.config.experimental.set_memory_growth(physical_devices[1], True)
#     tensorflow.config.experimental.set_memory_growth(physical_devices[2], True)
#     tensorflow.config.experimental.set_memory_growth(physical_devices[3], True)
    model = MobileNetV2(
        weights="imagenet",
        include_top=False,
        pooling="avg",
        input_shape=(224, 224, 3),
    )
    return model
from PIL import Image

# Use GPU wth
@ray.remote(num_gpus=1,num_cpus=5)
#@ray.remote
class generate_embedding_ray:
    
    def __init__(self,image_dict):
        self.image_dict = image_dict
        self.device_id = ray.get_gpu_ids() 
        print(self.device_id[0])
        self.model = create_keras_model(self.device_id[0])
    
    def generate_embeddings(self):
        y=[]
        from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
        from tensorflow.keras.preprocessing import image
        image_path =  self.image_dict['image_path']
        pil_image = Image.open(image_path)
        resized_image = pil_image.resize((224, 224))
        x = image.img_to_array(resized_image)
        #x = np.expand_dims(x, axis=0)
        y=np.append(x)
        y = preprocess_input(y)
        embeddings = self.model.predict(x).flatten().tolist()
#         embeddings_per_image.append(embeddings)
        return embeddings

output : 
(pid=1766620) 1
(pid=1766621) 2
(pid=1766620) 2021-04-02 03:36:14.810469: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
(pid=1766621) 2021-04-02 03:36:14.810471: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
(pid=1766620) 2021-04-02 03:36:16.163272: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
(pid=1766620) 2021-04-02 03:36:16.164073: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1
(pid=1766620) 2021-04-02 03:36:16.222705: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 0 with properties: 
(pid=1766620) pciBusID: 0000:88:00.0 name: Tesla V100-SXM2-32GB computeCapability: 7.0
(pid=1766620) coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 31.75GiB deviceMemoryBandwidth: 836.37GiB/s
(pid=1766620) 2021-04-02 03:36:16.222743: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
(pid=1766620) 2021-04-02 03:36:16.225993: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublas.so.11
(pid=1766620) 2021-04-02 03:36:16.226052: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublasLt.so.11
(pid=1766620) 2021-04-02 03:36:16.227347: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcufft.so.10
(pid=1766620) 2021-04-02 03:36:16.227634: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcurand.so.10
(pid=1766620) 2021-04-02 03:36:16.231153: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusolver.so.10
(pid=1766620) 2021-04-02 03:36:16.231903: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusparse.so.11
(pid=1766620) 2021-04-02 03:36:16.232051: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudnn.so.8
(pid=1766620) 2021-04-02 03:36:16.236750: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1862] Adding visible gpu devices: 0
(pid=1766621) 2021-04-02 03:36:16.271000: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
(pid=1766621) 2021-04-02 03:36:16.272049: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcuda.so.1
(pid=1766621) 2021-04-02 03:36:16.336261: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1720] Found device 0 with properties: 
(pid=1766621) pciBusID: 0000:b2:00.0 name: Tesla V100-SXM2-32GB computeCapability: 7.0
(pid=1766621) coreClock: 1.53GHz coreCount: 80 deviceMemorySize: 31.75GiB deviceMemoryBandwidth: 836.37GiB/s
(pid=1766621) 2021-04-02 03:36:16.336293: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudart.so.11.0
(pid=1766621) 2021-04-02 03:36:16.339520: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublas.so.11
(pid=1766621) 2021-04-02 03:36:16.339587: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcublasLt.so.11
(pid=1766621) 2021-04-02 03:36:16.340902: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcufft.so.10
(pid=1766621) 2021-04-02 03:36:16.341189: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcurand.so.10
(pid=1766621) 2021-04-02 03:36:16.344682: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusolver.so.10
(pid=1766621) 2021-04-02 03:36:16.345429: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcusparse.so.11
(pid=1766621) 2021-04-02 03:36:16.345574: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library libcudnn.so.8
(pid=1766621) 2021-04-02 03:36:16.353692: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1862] Adding visible gpu devices: 0
2021-04-02 03:36:21,250	ERROR worker.py:1018 -- Possible unhandled error from worker: ray::generate_embedding_ray.__init__() (pid=1766620, ip=172.20.6.92)
  File "python/ray/_raylet.pyx", line 484, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 438, in ray._raylet.execute_task.function_executor
  File "<ipython-input-49-496a8becdee1>", line 12, in __init__
  File "<ipython-input-48-bb48eca95898>", line 7, in create_keras_model
IndexError: list index out of range
2021-04-02 03:36:22,253	ERROR worker.py:1018 -- Possible unhandled error from worker: ray::generate_embedding_ray.__init__() (pid=1766621, ip=172.20.6.92)
  File "python/ray/_raylet.pyx", line 484, in ray._raylet.execute_task
  File "python/ray/_raylet.pyx", line 438, in ray._raylet.execute_task.function_executor
  File "<ipython-input-49-496a8becdee1>", line 12, in __init__
  File "<ipython-input-48-bb48eca95898>", line 7, in create_keras_model
IndexError: list index out of range

Can you try set GPU IDs to 0?