add support for a100

This commit is contained in:
2025-09-03 11:01:20 +08:00
parent f0425d03b4
commit fcf2001f24
3 changed files with 10 additions and 0 deletions

5
Dockerfile.a100 Normal file
View File

@@ -0,0 +1,5 @@
FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel
WORKDIR /workspace
RUN pip install sentence-transformers -i https://nexus.4pd.io/repository/pypi-all/simple
COPY . /workspace/

3
run_in_docker_a100.sh Executable file
View File

@@ -0,0 +1,3 @@
#! /usr/bin/env bash
image=harbor-contest.4pd.io/zhanghao/feature:a100-0.1
docker run -it -v /home/zhanghao/workspace:/workspace -v /mnt:/mnt $image bash

2
test.sh Executable file
View File

@@ -0,0 +1,2 @@
# python main.py dataset.json output.json /mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5 --device npu
python main.py dataset.json output.json /mnt/contest_ceph/zhanghao/models/BAAI/bge-large-zh-v1.5 --device cuda