E2E tests: Upgrade MNIST dependencies

sutaakar · sutaakar · commit ac559f2d2951 · 2024-12-16T14:24:35.000+01:00
diff --git a/test/e2e/mnist.py b/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)
diff --git a/test/e2e/mnist_pip_requirements.txt b/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

Original file line number	Diff line number	Diff line change
`@@ -72,8 +72,8 @@ def __init__(self, data_dir=PATH_DATASETS, hidden_size=64, learning_rate=2e-4):`
`72`	`72`	`nn.Linear(hidden_size, self.num_classes),`
`73`	`73`	`)`
`74`	`74`
`75`		`- self.val_accuracy = Accuracy()`
`76`		`- self.test_accuracy = Accuracy()`
	`75`	`+ self.val_accuracy = Accuracy(task="multiclass", num_classes=10)`
	`76`	`+ self.test_accuracy = Accuracy(task="multiclass", num_classes=10)`
`77`	`77`
`78`	`78`	`def forward(self, x):`
`79`	`79`	`x = self.model(x)`