Reduction on full image not required?

[mattclifford1]

https://github.com/Lightning-AI/metrics/blob/ff61c482e5157b43e647565fa0020a4ead6e9d61/torchmetrics/functional/image/ssim.py#L191-L194

To get the full SSIM image to be returned I have to change to reduction=None

Example: imports and dummy images

import torch
from torchmetrics import StructuralSimilarityIndexMeasure as SSIM

im1 = torch.rand([1, 1, 256, 256], generator=torch.manual_seed(42))
im2 = im1 * 0.75

Without reduction arg changed from default - get the mean SSIM twice:

metric = SSIM(return_full_image=True)
score = metric(im1, im2)
print(score)

>>> (tensor(0.9219), tensor(0.9219))

With reduction arg changed (desired behaviour of outputting mean SSIM and the full SSIM image)

metric = SSIM(return_full_image=True, reduction=None)
score = metric(im1, im2)
print(score)

>>> (tensor(0.9219), tensor([[[[0.9218, 0.9218, 0.9218,  ..., 0.9221, 0.9220, 0.9220],
          [0.9218, 0.9218, 0.9218,  ..., 0.9220, 0.9220, 0.9220],
          [0.9218, 0.9218, 0.9218,  ..., 0.9219, 0.9219, 0.9219],
          ...,
          [0.9219, 0.9219, 0.9219,  ..., 0.9219, 0.9219, 0.9219],
          [0.9219, 0.9219, 0.9219,  ..., 0.9220, 0.9220, 0.9220],
          [0.9219, 0.9219, 0.9219,  ..., 0.9220, 0.9220, 0.9221]]]]))

Does the return https://github.com/Lightning-AI/metrics/blob/ff61c482e5157b43e647565fa0020a4ead6e9d61/torchmetrics/functional/image/ssim.py#L191-L194

Need to be changed to:

elif return_full_image: 
     return (reduce(ssim_idx.reshape(ssim_idx.shape[0], -1).mean(-1), reduction), ssim_idx_full_image) 

or am I using the SSIM object incorrectly?

[github-actions[bot]]

Hi! thanks for your contribution!, great first issue!