train_job.py

from huggingface_hub import run_job


# docker pull pytorch/pytorch:2.8.0-cuda12.8-cudnn9-devel
job = run_job(
    image="pytorch/pytorch:2.7.1-cuda12.8-cudnn9-runtime",
    command=[
            "bash",
            "-c",
            "export CUDA_HOME=/usr/local/cuda && "
            "export PATH=$CUDA_HOME/bin:$PATH && "
            "export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH && "
            "apt-get update && apt-get install -y git && "
            "git clone https://huggingface.co/uav-vts/training_sem . && "
            "pwd && ls -al && "
            "pip install -r requirements.txt && "
            "python runner/train.py"
    ],
    flavor="a100-large",
    namespace="uav-vts"
    # token="hf_xxx"
)