Use newer PyTorch image for tests

sutaakar · sutaakar · commit cd5e0e5549a7 · 2024-12-16T19:26:50.000+01:00
diff --git a/support/environment.go b/support/environment.go
@@ -95,7 +95,7 @@ func GetRayTorchROCmImage() string {
 }
 
 func GetPyTorchImage() string {
-	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")
+	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:2.4.1-cuda11.8-cudnn9-runtime")
 }
 
 func GetCudaTrainingImage() string {

Original file line number	Diff line number	Diff line change
`@@ -95,7 +95,7 @@ func GetRayTorchROCmImage() string {`
`95`	`95`	`}`
`96`	`96`
`97`	`97`	`func GetPyTorchImage() string {`
`98`		`- return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")`
	`98`	`+ return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:2.4.1-cuda11.8-cudnn9-runtime")`
`99`	`99`	`}`
`100`	`100`
`101`	`101`	`func GetCudaTrainingImage() string {`