Скрипт Dockerized Tensorflow не видит графические процессоры

Question

Я использую docker контейнеров для обучения моделей глубокого обучения. Эти контейнеры Docker расположены на сервере Linux и обучаются там с помощью нескольких графических процессоров.

Проблема в том, что Tensorflow не распознает графические процессоры внутри контейнера. Контейнер Docker выглядит так:

FROM nvidia/cuda:10.2-runtime-ubuntu18.04
RUN apt-get update && apt-get install -y apt-utils
RUN apt-get install -y \
    git \
    pkg-config \
    python3-pip \
    python3.6 \
    nano \ 
    wget \
    yasm


FROM python:3.6
COPY requirements.txt ./

# Here tensorflow-gpu == 2.1 is installed
RUN pip install --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt

COPY . /
ENTRYPOINT ["python", "./main.py"]

Если вы сейчас закомментируете строки python specifici c из файла docker и замените их на CMD ["nvidia-smi"] вы можете видеть, что графические процессоры внутри контейнера видны. Теперь единственный вопрос, который я должен задать себе, - как Tensorflow может обнаруживать графические процессоры.

В коде Python графические процессоры включены следующим образом:

physical_devices = tf.config.experimental.list_physical_devices('GPU')
    for physical_device in physical_devices:
        tf.config.experimental.set_memory_growth(physical_device, True)