From 3b69661f056d382eda4d45a4230673e13a520590 Mon Sep 17 00:00:00 2001 From: Adam Dangoor Date: Sun, 4 Feb 2024 18:35:40 +0000 Subject: [PATCH] Fix return type of piq.ssim This function can return a single tensor if the parameter `full` is set to `False` (the default). --- piq/ssim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/piq/ssim.py b/piq/ssim.py index 74814ce5..48619172 100644 --- a/piq/ssim.py +++ b/piq/ssim.py @@ -19,7 +19,7 @@ def ssim(x: torch.Tensor, y: torch.Tensor, kernel_size: int = 11, kernel_sigma: float = 1.5, data_range: Union[int, float] = 1., reduction: str = 'mean', full: bool = False, - downsample: bool = True, k1: float = 0.01, k2: float = 0.03) -> List[torch.Tensor]: + downsample: bool = True, k1: float = 0.01, k2: float = 0.03) -> List[torch.Tensor] | torch.Tensor: r"""Interface of Structural Similarity (SSIM) index. Inputs supposed to be in range ``[0, data_range]``. To match performance with skimage and tensorflow set ``'downsample' = True``.