diff --git a/pyrad/graphs/instant_ngp.py b/pyrad/graphs/instant_ngp.py
index 5ff49aa01..b094521b1 100644
--- a/pyrad/graphs/instant_ngp.py
+++ b/pyrad/graphs/instant_ngp.py
@@ -20,8 +20,9 @@ from typing import Dict, List
 
 import torch
 from torch.nn import Parameter
-from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.functional import structural_similarity_index_measure
 
 from pyrad.fields.modules.field_heads import FieldHeadNames
 from pyrad.fields.instant_ngp_field import field_implementation_to_class
@@ -81,7 +82,7 @@ class NGPGraph(Graph):
 
         # metrics
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
-        self.ssim = StructuralSimilarityIndexMeasure()
+        self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
diff --git a/pyrad/graphs/mipnerf.py b/pyrad/graphs/mipnerf.py
index b8b6958a8..9f3f5da9c 100644
--- a/pyrad/graphs/mipnerf.py
+++ b/pyrad/graphs/mipnerf.py
@@ -21,8 +21,9 @@ from typing import Dict, List
 
 import torch
 from torch.nn import Parameter
-from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.functional import structural_similarity_index_measure
 
 from pyrad.fields.modules.encoding import NeRFEncoding
 from pyrad.fields.modules.field_heads import FieldHeadNames
@@ -85,7 +86,7 @@ class MipNerfGraph(Graph):
 
         # metrics
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
-        self.ssim = StructuralSimilarityIndexMeasure()
+        self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
diff --git a/pyrad/graphs/mipnerf_360.py b/pyrad/graphs/mipnerf_360.py
index 35343057a..147094309 100644
--- a/pyrad/graphs/mipnerf_360.py
+++ b/pyrad/graphs/mipnerf_360.py
@@ -21,8 +21,9 @@ from typing import Dict, List
 
 import torch
 from torch.nn import Parameter
-from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.functional import structural_similarity_index_measure
 
 from pyrad.fields.modules.encoding import NeRFEncoding
 from pyrad.fields.modules.field_heads import FieldHeadNames
@@ -89,7 +90,7 @@ class MipNerf360Graph(Graph):
 
         # metrics
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
-        self.ssim = StructuralSimilarityIndexMeasure()
+        self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
diff --git a/pyrad/graphs/nerfw.py b/pyrad/graphs/nerfw.py
index 7f6b01969..981712ab0 100644
--- a/pyrad/graphs/nerfw.py
+++ b/pyrad/graphs/nerfw.py
@@ -17,8 +17,9 @@ NeRF-W (NeRF in the wild) implementation.
 """
 
 import torch
-from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.functional import structural_similarity_index_measure
 
 from pyrad.fields.modules.encoding import NeRFEncoding
 from pyrad.fields.modules.field_heads import FieldHeadNames
@@ -92,7 +93,7 @@ class NerfWGraph(Graph):
 
         # metrics
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
-        self.ssim = StructuralSimilarityIndexMeasure()
+        self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
     def get_param_groups(self):
diff --git a/pyrad/graphs/vanilla_nerf.py b/pyrad/graphs/vanilla_nerf.py
index db1782d80..54bf56bb2 100644
--- a/pyrad/graphs/vanilla_nerf.py
+++ b/pyrad/graphs/vanilla_nerf.py
@@ -21,8 +21,9 @@ from typing import Dict, List
 
 import torch
 from torch.nn import Parameter
-from torchmetrics import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
+from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
+from torchmetrics.functional import structural_similarity_index_measure
 
 from pyrad.fields.modules.encoding import NeRFEncoding
 from pyrad.fields.modules.field_heads import FieldHeadNames
@@ -94,7 +95,7 @@ class NeRFGraph(Graph):
 
         # metrics
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
-        self.ssim = StructuralSimilarityIndexMeasure()
+        self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity()
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]: