Mirrors/modelscope
1
0
Fork 0
mirror of https://github.com/modelscope/modelscope.git synced 2025-12-16 16:27:45 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix tensorflow cannot load gpu (#1456)

Browse Source
This commit is contained in:
tastelikefeet
2025-08-11 11:04:17 +08:00
committed by GitHub
parent 055496c597
commit 169f523082
1 changed files with 4 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
docker/build_image.py
View File

@@ -204,11 +204,15 @@ class GPUImageBuilder(Builder):
def generate_dockerfile(self) -> str: def generate_dockerfile(self) -> str:
meta_file = './docker/install.sh' meta_file = './docker/install.sh'
# pushd ~ popd is to solve the tf cannot use gpu problem.
extra_content = """ extra_content = """
RUN pip install tf-keras==2.16.0 --no-dependencies && \ RUN pip install tf-keras==2.16.0 --no-dependencies && \
pip install --no-cache-dir torchsde jupyterlab torchmetrics==0.11.4 basicsr pynvml shortuuid && \ 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" \ 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' 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 = ( version_args = (