diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go
index 7c5d3adedc..b2ac705206 100644
--- a/pkg/dockerfile/generator.go
+++ b/pkg/dockerfile/generator.go
@@ -97,7 +97,7 @@ func (g *Generator) GenerateBase() (string, error) {
 		return "", err
 	}
 	installPython := ""
-	if g.Config.Build.GPU {
+	if g.Config.Build.GPU && g.useCudaBaseImage {
 		installPython, err = g.installPythonCUDA()
 		if err != nil {
 			return "", err