From 1dc9502dd9fcb0184d70166839490a2a91ce7b96 Mon Sep 17 00:00:00 2001
From: tastelikefeet <58414341+tastelikefeet@users.noreply.github.com>
Date: Mon, 11 Aug 2025 11:04:17 +0800
Subject: [PATCH] fix tensorflow cannot load gpu (#1456)

(cherry picked from commit 169f52308268e60bddb4e19a2785f7cb473c0f39)
---
 docker/build_image.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/docker/build_image.py b/docker/build_image.py
index bfef7be5..e1b6ca14 100644
--- a/docker/build_image.py
+++ b/docker/build_image.py
@@ -204,11 +204,15 @@ class GPUImageBuilder(Builder):
 
     def generate_dockerfile(self) -> str:
         meta_file = './docker/install.sh'
+        # pushd ~ popd is to solve the tf cannot use gpu problem.
         extra_content = """
 RUN pip install tf-keras==2.16.0 --no-dependencies && \
     pip install --no-cache-dir torchsde jupyterlab torchmetrics==0.11.4 basicsr pynvml shortuuid && \
     CUDA_HOME=/usr/local/cuda TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6 8.9 9.0" \
         pip install --no-cache-dir  'git+https://github.com/facebookresearch/detectron2.git'
+RUN pushd $(dirname $(python -c 'print(__import__("tensorflow").__file__)'))  && \
+    ln -svf ../nvidia/*/lib/*.so* .  && \
+    popd
 """
 
         version_args = (