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

# 安装必要工具
RUN apt-get update && apt-get install -y git curl wget nano openssh-server \
    && rm -rf /var/lib/apt/lists/*

# pip依赖
RUN pip install --no-cache-dir tqdm numpy requests

# 创建 vscode-server 需要的目录
RUN mkdir -p /root/.vscode-server /root/.vscode-server/data/Machine

# 设置工作目录
WORKDIR /workspace

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