From 1ab3ea21cee46fbb05b8e5f6f5ceea5c5e212518 Mon Sep 17 00:00:00 2001
From: Valerio Mariani <valeriomariani52@gmail.com>
Date: Wed, 30 Oct 2024 21:53:32 +0100
Subject: [PATCH] make base image parametric.

default `BASE_IMAGE` is nvcr.io/nvidia/pytorch:24.02-py3 for retro-compatibility


Former-commit-id: d5b2def7c029d6f808a97cb062340810e0cd075d
---
 docker/docker-cuda/Dockerfile | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/docker/docker-cuda/Dockerfile b/docker/docker-cuda/Dockerfile
index d03ece88..f4c12462 100644
--- a/docker/docker-cuda/Dockerfile
+++ b/docker/docker-cuda/Dockerfile
@@ -1,6 +1,5 @@
-# Use the NVIDIA official image with PyTorch 2.3.0
-# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-02.html
-FROM nvcr.io/nvidia/pytorch:24.02-py3
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:24.02-py3
+FROM ${BASE_IMAGE}
 
 # Define environments
 ENV MAX_JOBS=4