diff --git a/pkg/config/config.go b/pkg/config/config.go
index 95bf7e4656..b971a03ab5 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -31,7 +31,9 @@ type RunItem struct {
 }
 
 type Build struct {
-	GPU                bool      `json:"gpu,omitempty" yaml:"gpu"`
+	GPU bool `json:"gpu,omitempty" yaml:"gpu"`
+	// UseSystemCuda? something about jax?
+	UseCudaBaseImage   bool      `json:"use_cuda_base_image,omitempty" yaml:"use_cuda_base_image"`
 	PythonVersion      string    `json:"python_version,omitempty" yaml:"python_version"`
 	PythonRequirements string    `json:"python_requirements,omitempty" yaml:"python_requirements"`
 	PythonPackages     []string  `json:"python_packages,omitempty" yaml:"python_packages"` // Deprecated, but included for backwards compatibility
@@ -59,8 +61,9 @@ type Config struct {
 func DefaultConfig() *Config {
 	return &Config{
 		Build: &Build{
-			GPU:           false,
-			PythonVersion: "3.8",
+			GPU:              false,
+			PythonVersion:    "3.8",
+			UseCudaBaseImage: true,
 		},
 	}
 }
diff --git a/pkg/config/data/config_schema_v1.0.json b/pkg/config/data/config_schema_v1.0.json
index 7c7a48468f..1e2bc7d499 100644
--- a/pkg/config/data/config_schema_v1.0.json
+++ b/pkg/config/data/config_schema_v1.0.json
@@ -24,6 +24,11 @@
           "type": "boolean",
           "description": "Enable GPUs for this model. When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used, and Cog will automatically figure out what versions of CUDA and cuDNN to use based on the version of Python, PyTorch, and Tensorflow that you are using."
         },
+        "use_cuda_base_image": {
+          "$id": "#/properties/build/properties/use_cuda_base_image",
+          "type": "boolean",
+          "description": "When enabled, the [nvidia-docker](https://github.com/NVIDIA/nvidia-docker) base image will be used. This is not normally necessary for torch, but may be necessary for other frameworks."
+        },
         "python_version": {
           "$id": "#/properties/build/properties/python_version",
           "type": ["string", "number"],
diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go
index 99b1f93eb5..c603cf04ba 100644
--- a/pkg/dockerfile/generator.go
+++ b/pkg/dockerfile/generator.go
@@ -158,7 +158,7 @@ func (g *Generator) Generate(imageName string) (weightsBase string, dockerfile s
 		return "", "", "", err
 	}
 	installPython := ""
-	if g.Config.Build.GPU {
+	if g.Config.Build.GPU && g.Config.Build.UseCudaBaseImage {
 		installPython, err = g.installPythonCUDA()
 		if err != nil {
 			return "", "", "", err
@@ -245,10 +245,10 @@ func (g *Generator) Cleanup() error {
 }
 
 func (g *Generator) baseImage() (string, error) {
-	if g.Config.Build.GPU {
+	if g.Config.Build.GPU && g.Config.Build.UseCudaBaseImage {
 		return g.Config.CUDABaseImageTag()
 	}
-	return "python:" + g.Config.Build.PythonVersion, nil
+	return "python:" + g.Config.Build.PythonVersion + "-slim", nil
 }
 
 func (g *Generator) preamble() string {
diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go
index 350b132028..693b6a1211 100644
--- a/pkg/dockerfile/generator_test.go
+++ b/pkg/dockerfile/generator_test.go
@@ -81,7 +81,7 @@ predict: predict.py:Predictor
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -152,7 +152,7 @@ predict: predict.py:Predictor
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -248,7 +248,7 @@ build:
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -428,7 +428,7 @@ predict: predict.py:Predictor
 
 	expected := `#syntax=docker/dockerfile:1.4
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
diff --git a/test-integration/test_integration/test_build.py b/test-integration/test_integration/test_build.py
index 67bd67ff4d..0260c7a0f8 100644
--- a/test-integration/test_integration/test_build.py
+++ b/test-integration/test_integration/test_build.py
@@ -31,7 +31,7 @@ def test_build_without_predictor(docker_image):
     # Deprecated. Remove for 1.0.
     assert len(labels["org.cogmodel.cog_version"]) > 0
     assert json.loads(labels["org.cogmodel.config"]) == {
-        "build": {"python_version": "3.8"}
+        "build": {"python_version": "3.8", "use_cuda_base_image": True}
     }
     assert "org.cogmodel.openapi_schema" not in labels
 
@@ -155,6 +155,7 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image):
             "gpu": True,
             "cuda": "11.8",
             "cudnn": "8",
+            "use_cuda_base_image": True,
         }
     }
     assert "run.cog.openapi_schema" not in labels
@@ -167,6 +168,7 @@ def test_build_gpu_model_on_cpu(tmpdir, docker_image):
             "gpu": True,
             "cuda": "11.8",
             "cudnn": "8",
+            "use_cuda_base_image": True,
         }
     }
     assert "org.cogmodel.openapi_schema" not in labels