fix: remove deprecated pynvml, remove torchmetrics restrictions (#566)

* fix: remove deprecated pynvml, remove torchmetrics restrictions

Author: begumcig

pyproject.toml

@@ -113,7 +113,7 @@ classifiers = [
 dependencies = [
     "torch>=2.7.0",
     "torchvision>=0.22.0",
-    "torchmetrics[image]==1.7.4",
+    "torchmetrics[image]",
     "requests>=2.31.0",
     "transformers<5.0.0",
     "pytorch-lightning",
@@ -134,7 +134,7 @@ dependencies = [
     "opentelemetry-sdk>=1.30.0",
     "opentelemetry-exporter-otlp>=1.29.0",
     "codecarbon",
-    "pynvml",
+    "nvidia-ml-py",
     "thop",
     "timm",
     "bitsandbytes; sys_platform != 'darwin' or platform_machine != 'arm64'",

src/pruna/evaluation/metrics/metric_torch.py

@@ -179,7 +179,7 @@ class TorchMetrics(Enum):
     clip_score = (partial(CLIPScore), None, "y_x")
     precision = (partial(Precision), None, "y_gt")
     recall = (partial(Recall), None, "y_gt")
-    psnr = (partial(PeakSignalNoiseRatio), None, "pairwise_y_gt")
+    psnr = (partial(PeakSignalNoiseRatio, data_range=255.0), None, "pairwise_y_gt")
     ssim = (partial(StructuralSimilarityIndexMeasure), ssim_update, "pairwise_y_gt")
     msssim = (partial(MultiScaleStructuralSimilarityIndexMeasure), ssim_update, "pairwise_y_gt")
     lpips = (partial(LearnedPerceptualImagePatchSimilarity), lpips_update, "pairwise_y_gt")