CANN:8.3.rc1
HDK:25.2.0
python：3.11.13
torch:2.1.0
torch_npu:2.1.0.post17
cloudpickle:1.6.1

docker run -it  -u root \
--net=host --shm-size=5g \
--device=/dev/davinci_manager \
--device=/dev/devmm_svm \
--device=/dev/hisi_hdc \
--device=/dev/davinci4 \
--device=/dev/davinci5 \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
-v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /usr/share/zoneinfo/Asia/Shanghai:/etc/localtime \
-v /home:/home/ \
--name chemprop_test \
--entrypoint=/bin/bash \
-it swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11

wget https://download.pytorch.org/whl/cpu/torch-2.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
wget https://gitcode.com/Ascend/pytorch/releases/download/v7.2.0-pytorch2.1.0/torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

pip3 install torch-2.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
pip3 install torch_npu-2.1.0.post17-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

apt update && apt install -y libsm6 libxext6

pip3 install astartes aimsim configargparse lightning>=2.0 scikit-learn==1.1.3 descriptastorus rich cloudpickle ml-dtypes tornado numpy==1.26.4

git clone https://atomgit.com/AI4Science/chemprop.git
cd chemprop && pip install -e .

source /usr/local/Ascend/ascend-toolkit/set_env.sh
python3 train.py --data_path tests/data/classification.csv --dataset_type classification --save_dir classification_checkpoints