experiment with nvcr pytorch image for torch 2.5.1

2024-12-06 11:07:27 -05:00
parent 5bef19064b
commit 684b543aa1
3 changed files with 15 additions and 8 deletions
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,5 +1,6 @@
+ARG BASE_IMAGE=axolotlai/axolotl-base
 ARG BASE_TAG=main-base
-FROM axolotlai/axolotl-base:$BASE_TAG
+FROM $BASE_IMAGE:$BASE_TAG

 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6+PTX"
 ARG AXOLOTL_EXTRAS=""