FROM nvcr.io/nvidia/pytorch:24.02-py3 ENV TORCH_CUDA_ARCH_LIST "7.0;7.5;8.0;8.6;8.7;8.9;9.0" ENV MMCV_WITH_OPS 1 ENV FORCE_CUDA 1 RUN python3 -m pip install --upgrade pip # Install, then uninstall to get only the deps COPY . ./megatron-energon RUN pip install -e ./megatron-energon && pip uninstall -y megatron-energon && rm -rf ./megatron-energon
