Jelajahi Sumber

build(docker): add torch and torchvision installation

- Add pip install command for torch and torchvision
- Specify version2.3.1 for both packages
- Use Aliyun mirror for faster download
myhloli 7 bulan lalu
induk
melakukan
5ddd6799aa
1 mengubah file dengan 1 tambahan dan 0 penghapusan
  1. 1 0
      docker/ascend_npu/Dockerfile

+ 1 - 0
docker/ascend_npu/Dockerfile

@@ -35,6 +35,7 @@ RUN /bin/bash -c "wget https://gcore.jsdelivr.net/gh/opendatalab/MinerU@master/m
     cp magic-pdf.template.json /root/magic-pdf.json && \
     source /opt/mineru_venv/bin/activate && \
     pip3 install --upgrade pip -i https://mirrors.aliyun.com/pypi/simple && \
+    pip3 install torch==2.3.1 torchvision==0.18.1 -i https://mirrors.aliyun.com/pypi/simple && \
     pip3 install -U magic-pdf[full] -i https://mirrors.aliyun.com/pypi/simple && \
     wget https://gitee.com/ascend/pytorch/releases/download/v6.0.rc2-pytorch2.3.1/torch_npu-2.3.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl && \
     pip3 install torch_npu-2.3.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl"