E2E tests: Upgrade MNIST test script dependencies

Author: sutaakar

go.mod

@@ -11,7 +11,7 @@ require (
 	github.com/openshift/api v0.0.0-20230823114715-5fdd7511b790
 	github.com/openshift/client-go v0.0.0-20221019143426-16aed247da5c
 	github.com/project-codeflare/appwrapper v0.27.0
-	github.com/project-codeflare/codeflare-common v0.0.0-20240930133152-11fd6e3be6b3
+	github.com/project-codeflare/codeflare-common v0.0.0-20241216092650-1f9ad0399709
 	github.com/ray-project/kuberay/ray-operator v1.2.1
 	go.uber.org/zap v1.27.0
 	golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8

go.sum

@@ -226,8 +226,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/project-codeflare/appwrapper v0.27.0 h1:WiWw0Hi4rEXuFIEpm8nq1UqJHgVB6YtGcWzRrhRUTyE=
 github.com/project-codeflare/appwrapper v0.27.0/go.mod h1:7FpO90DLv0BAq4rwZtXKS9aRRfkR9RvXsj3pgYF0HtQ=
-github.com/project-codeflare/codeflare-common v0.0.0-20240930133152-11fd6e3be6b3 h1:Eupu9yxaGTddtoxb9SjrYJlokRHEYU5NNVRQmdXSNVs=
-github.com/project-codeflare/codeflare-common v0.0.0-20240930133152-11fd6e3be6b3/go.mod h1:v7XKwaDoCspsHQlWJNarO7gOpR+iumSS+c1bWs3kJOI=
+github.com/project-codeflare/codeflare-common v0.0.0-20241216092650-1f9ad0399709 h1:/cjC/eR9SXqP/r4x6FD1BcvAyqhmfN+jl0W53bC+1EI=
+github.com/project-codeflare/codeflare-common v0.0.0-20241216092650-1f9ad0399709/go.mod h1:DPSv5khRiRDFUD43SF8da+MrVQTWmxNhuKJmwSLOyO0=
 github.com/prometheus/client_golang v1.20.4 h1:Tgh3Yr67PaOv/uTqloMsCEdeuFTatm5zIq5+qNN23vI=
 github.com/prometheus/client_golang v1.20.4/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=

test/e2e/mnist.py

@@ -72,8 +72,8 @@ def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):
             nn.Linear(hidden_size, self.num_classes),
         )
 
-        self.val_accuracy = Accuracy()
-        self.test_accuracy = Accuracy()
+        self.val_accuracy = Accuracy(task="multiclass", num_classes=10)
+        self.test_accuracy = Accuracy(task="multiclass", num_classes=10)
 
     def forward(self, x):
         x = self.model(x)

test/e2e/mnist_pip_requirements.txt

@@ -1,3 +1,3 @@
-pytorch_lightning==1.9.5
-torchmetrics==0.9.1
-torchvision==0.12.0
+pytorch_lightning==2.4.0
+torchmetrics==1.6.0
+torchvision==0.20.1