Add test for using Torch & Tensorflow together (#1023)

* Failing test for using Torch & Tensorflow together

Signed-off-by: Ben Firshman <ben@firshman.com>

* Update to latest versions of CUDA, PyTorch and TensorFlow

---------

Signed-off-by: Ben Firshman <ben@firshman.com>
Co-authored-by: Mattt Zmuda <mattt@replicate.com>

pkg/config/config_test.go

@@ -470,6 +470,31 @@ foo==1.0.0`
 	require.NotContains(t, requirements, "tensorflow_gpu")
 }
 
+func TestPythonPackagesBothTorchAndTensorflow(t *testing.T) {
+	config := &Config{
+		Build: &Build{
+			GPU:           true,
+			PythonVersion: "3.11",
+			PythonPackages: []string{
+				"tensorflow==2.16.1",
+				"torch==2.3.1",
+			},
+			CUDA: "12.3",
+		},
+	}
+	err := config.ValidateAndComplete("")
+	require.NoError(t, err)
+	require.Equal(t, "12.3", config.Build.CUDA)
+	require.Equal(t, "8", config.Build.CuDNN)
+
+	requirements, err := config.PythonRequirementsForArch("", "", []string{})
+	require.NoError(t, err)
+	expected := `--extra-index-url https://download.pytorch.org/whl/cu121
+tensorflow==2.16.1
+torch==2.3.1+cu121`
+	require.Equal(t, expected, requirements)
+}
+
 func TestCUDABaseImageTag(t *testing.T) {
 	config := &Config{
 		Build: &Build{