FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

# 基础工具
RUN apt-get update && apt-get install -y git curl wget nano && rm -rf /var/lib/apt/lists/*

# tqdm用于进度条
RUN pip install --no-cache-dir tqdm numpy requests

# 设置工作目录
WORKDIR /workspace

# 验证CUDA
RUN python3 -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"