From cd5e0e5549a766231b493be9ab5e6f0bfe04b165 Mon Sep 17 00:00:00 2001
From: Karel Suta <ksuta@redhat.com>
Date: Mon, 16 Dec 2024 19:26:36 +0100
Subject: [PATCH] Use newer PyTorch image for tests

---
 support/environment.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/support/environment.go b/support/environment.go
index fb4cd6a..abb2f1d 100644
--- 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 {