We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
The sdr metric in TM sometime gives NaN for some input
## 🐛 Bug

This issue is related with `fast-bss-eval`'s torch version, see https://github.com/fakufaku/fast_bss_eval/issues/5

### To Reproduce

```
import numpy as np
import torch

x = np.load('debug.npz')
preds = torch.tensor(x['preds'])
target = torch.tensor(x['target'])
print(preds.shape, target.shape)

from torchmetrics.functional.audio import signal_distortion_ratio
sdr = signal_distortion_ratio(preds, target)
print(sdr)

from mir_eval.separation import bss_eval_sources
sdr, _, _, _ = bss_eval_sources(target.numpy(), preds.numpy(), False)
print(sdr)
```

outputs:
```
torch.Size([2, 64000]) torch.Size([2, 64000])
tensor([-2.6815,     nan])
[-2.68156071 44.58523729]
```

unzip [data.zip](https://github.com/PyTorchLightning/metrics/files/8309638/data.zip) to get the `debug.npz`


<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

#### Code sample

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

the results given by signal_distortion_ratio is close to the one given by `mir_eval`
<!-- A clear and concise description of what you expected to happen. -->

### Environment

- OS (e.g., Linux):
- Python & PyTorch Version (e.g., 1.0):
- How you installed PyTorch (`conda`, `pip`, build command if you used source):
- Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchmetrics/functional/audio/sdr.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from typing import Optional
16 
17 import torch
18 
19 from torchmetrics.utilities.imports import _FAST_BSS_EVAL_AVAILABLE, _TORCH_GREATER_EQUAL_1_8
20 
21 if _FAST_BSS_EVAL_AVAILABLE:
22     if _TORCH_GREATER_EQUAL_1_8:
23         from fast_bss_eval.torch.cgd import toeplitz_conjugate_gradient
24         from fast_bss_eval.torch.helpers import _normalize
25         from fast_bss_eval.torch.linalg import toeplitz
26         from fast_bss_eval.torch.metrics import compute_stats
27 
28         solve = torch.linalg.solve
29     else:
30         import numpy
31         from fast_bss_eval.numpy.cgd import toeplitz_conjugate_gradient
32         from fast_bss_eval.numpy.helpers import _normalize
33         from fast_bss_eval.numpy.linalg import toeplitz
34         from fast_bss_eval.numpy.metrics import compute_stats
35 
36         solve = numpy.linalg.solve
37 else:
38     toeplitz = None
39     toeplitz_conjugate_gradient = None
40     compute_stats = None
41     _normalize = None
42     __doctest_skip__ = ["signal_distortion_ratio"]
43 
44 from torch import Tensor
45 
46 from torchmetrics.utilities import rank_zero_warn
47 from torchmetrics.utilities.checks import _check_same_shape
48 
49 
50 def signal_distortion_ratio(
51     preds: Tensor,
52     target: Tensor,
53     use_cg_iter: Optional[int] = None,
54     filter_length: int = 512,
55     zero_mean: bool = False,
56     load_diag: Optional[float] = None,
57 ) -> Tensor:
58     r"""Signal to Distortion Ratio (SDR) [1,2,3]
59 
60     Args:
61         preds:
62             shape ``[..., time]``
63         target:
64             shape ``[..., time]``
65         use_cg_iter:
66             If provided, an iterative method is used to solve for the distortion
67             filter coefficients instead of direct Gaussian elimination.
68             This can speed up the computation of the metrics in case the filters
69             are long. Using a value of 10 here has been shown to provide
70             good accuracy in most cases and is sufficient when using this
71             loss to train neural separation networks.
72         filter_length:
73             The length of the distortion filter allowed
74         zero_mean:
75             When set to True, the mean of all signals is subtracted prior to computation of the metrics
76         load_diag:
77             If provided, this small value is added to the diagonal coefficients of
78             the system metrics when solving for the filter coefficients.
79             This can help stabilize the metric in the case where some of the reference
80             signals may sometimes be zero
81 
82     Raises:
83         ModuleNotFoundError:
84             If ``fast-bss-eval`` package is not installed
85 
86     Returns:
87         sdr value of shape ``[...]``
88 
89     Example:
90 
91         >>> from torchmetrics.functional.audio import signal_distortion_ratio
92         >>> import torch
93         >>> g = torch.manual_seed(1)
94         >>> preds = torch.randn(8000)
95         >>> target = torch.randn(8000)
96         >>> signal_distortion_ratio(preds, target)
97         tensor(-12.0589)
98         >>> # use with permutation_invariant_training
99         >>> from torchmetrics.functional.audio import permutation_invariant_training
100         >>> preds = torch.randn(4, 2, 8000)  # [batch, spk, time]
101         >>> target = torch.randn(4, 2, 8000)
102         >>> best_metric, best_perm = permutation_invariant_training(preds, target, signal_distortion_ratio, 'max')
103         >>> best_metric
104         tensor([-11.6375, -11.4358, -11.7148, -11.6325])
105         >>> best_perm
106         tensor([[1, 0],
107                 [0, 1],
108                 [1, 0],
109                 [0, 1]])
110 
111     .. note::
112        1. when pytorch<1.8.0, numpy will be used to calculate this metric, which causes ``sdr`` to be
113             non-differentiable and slower to calculate
114 
115        2. using this metrics requires you to have ``fast-bss-eval`` install. Either install as ``pip install
116           torchmetrics[audio]`` or ``pip install fast-bss-eval``
117 
118        3. preds and target need to have the same dtype, otherwise target will be converted to preds' dtype
119 
120 
121     References:
122         [1] Vincent, E., Gribonval, R., & Fevotte, C. (2006). Performance measurement in blind audio source separation.
123         IEEE Transactions on Audio, Speech and Language Processing, 14(4), 1462–1469.
124 
125         [2] Scheibler, R. (2021). SDR -- Medium Rare with Fast Computations.
126 
127         [3] https://github.com/fakufaku/fast_bss_eval
128     """
129     if not _FAST_BSS_EVAL_AVAILABLE:
130         raise ModuleNotFoundError(
131             "SDR metric requires that `fast-bss-eval` is installed."
132             " Either install as `pip install torchmetrics[audio]` or `pip install fast-bss-eval`."
133         )
134     _check_same_shape(preds, target)
135 
136     if not preds.dtype.is_floating_point:
137         preds = preds.float()  # for torch.norm
138 
139     # half precision support
140     if preds.dtype == torch.float16:
141         preds = preds.to(torch.float32)
142 
143     if preds.dtype != target.dtype:  # for torch.linalg.solve
144         target = target.to(preds.dtype)
145 
146     if zero_mean:
147         preds = preds - preds.mean(dim=-1, keepdim=True)
148         target = target - target.mean(dim=-1, keepdim=True)
149 
150     # normalize along time-axis
151     if not _TORCH_GREATER_EQUAL_1_8:
152         # use numpy if torch<1.8
153         rank_zero_warn(
154             "Pytorch is under 1.8, thus SDR numpy version is used."
155             "For better performance and differentiability, you should change to Pytorch v1.8 or above."
156         )
157         device = preds.device
158         preds = preds.detach().cpu().numpy()
159         target = target.detach().cpu().numpy()
160 
161         preds = _normalize(preds, axis=-1)
162         target = _normalize(target, axis=-1)
163     else:
164         preds = _normalize(preds, dim=-1)
165         target = _normalize(target, dim=-1)
166 
167     # solve for the optimal filter
168     # compute auto-correlation and cross-correlation
169     acf, xcorr = compute_stats(target, preds, length=filter_length, pairwise=False)
170 
171     if load_diag is not None:
172         # the diagonal factor of the Toeplitz matrix is the first
173         # coefficient of the acf
174         acf[..., 0] += load_diag
175 
176     if use_cg_iter is not None:
177         # use preconditioned conjugate gradient
178         sol = toeplitz_conjugate_gradient(acf, xcorr, n_iter=use_cg_iter)
179     else:
180         # regular matrix solver
181         r_mat = toeplitz(acf)
182         sol = solve(r_mat, xcorr)
183 
184     # to tensor if torch<1.8
185     if not _TORCH_GREATER_EQUAL_1_8:
186         sol = torch.tensor(sol, device=device)
187         xcorr = torch.tensor(xcorr, device=device)
188 
189     # compute the coherence
190     coh = torch.einsum("...l,...l->...", xcorr, sol)
191 
192     # transform to decibels
193     ratio = coh / (1 - coh)
194     val = 10.0 * torch.log10(ratio)
195     return val
196 
197 
198 def scale_invariant_signal_distortion_ratio(preds: Tensor, target: Tensor, zero_mean: bool = False) -> Tensor:
199     """Calculates Scale-invariant signal-to-distortion ratio (SI-SDR) metric. The SI-SDR value is in general
200     considered an overall measure of how good a source sound.
201 
202     Args:
203         preds:
204             shape ``[...,time]``
205         target:
206             shape ``[...,time]``
207         zero_mean:
208             If to zero mean target and preds or not
209 
210     Returns:
211         si-sdr value of shape [...]
212 
213     Example:
214         >>> from torchmetrics.functional.audio import scale_invariant_signal_distortion_ratio
215         >>> target = torch.tensor([3.0, -0.5, 2.0, 7.0])
216         >>> preds = torch.tensor([2.5, 0.0, 2.0, 8.0])
217         >>> scale_invariant_signal_distortion_ratio(preds, target)
218         tensor(18.4030)
219 
220     References:
221         [1] Le Roux, Jonathan, et al. "SDR half-baked or well done." IEEE International Conference on Acoustics, Speech
222         and Signal Processing (ICASSP) 2019.
223     """
224     _check_same_shape(preds, target)
225     eps = torch.finfo(preds.dtype).eps
226 
227     if zero_mean:
228         target = target - torch.mean(target, dim=-1, keepdim=True)
229         preds = preds - torch.mean(preds, dim=-1, keepdim=True)
230 
231     alpha = (torch.sum(preds * target, dim=-1, keepdim=True) + eps) / (
232         torch.sum(target ** 2, dim=-1, keepdim=True) + eps
233     )
234     target_scaled = alpha * target
235 
236     noise = target_scaled - preds
237 
238     val = (torch.sum(target_scaled ** 2, dim=-1) + eps) / (torch.sum(noise ** 2, dim=-1) + eps)
239     val = 10 * torch.log10(val)
240 
241     return val
242
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/torchmetrics/functional/audio/sdr.py b/torchmetrics/functional/audio/sdr.py
--- a/torchmetrics/functional/audio/sdr.py
+++ b/torchmetrics/functional/audio/sdr.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from typing import Optional
 
 import torch
@@ -192,6 +193,15 @@
     # transform to decibels
     ratio = coh / (1 - coh)
     val = 10.0 * torch.log10(ratio)
+
+    # recompute sdr in float64 if val is NaN or Inf
+    if (torch.isnan(val).any() or torch.isinf(val).any()) and preds.dtype != torch.float64:
+        warnings.warn(
+            "Detected `nan` or `inf` value in computed metric, retrying computation in double precision",
+            UserWarning,
+        )
+        val = signal_distortion_ratio(preds.double(), target.double(), use_cg_iter, filter_length, zero_mean, load_diag)
+
     return val

{"golden_diff": "diff --git a/torchmetrics/functional/audio/sdr.py b/torchmetrics/functional/audio/sdr.py\n--- a/torchmetrics/functional/audio/sdr.py\n+++ b/torchmetrics/functional/audio/sdr.py\n@@ -12,6 +12,7 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n \n+import warnings\n from typing import Optional\n \n import torch\n@@ -192,6 +193,15 @@\n     # transform to decibels\n     ratio = coh / (1 - coh)\n     val = 10.0 * torch.log10(ratio)\n+\n+    # recompute sdr in float64 if val is NaN or Inf\n+    if (torch.isnan(val).any() or torch.isinf(val).any()) and preds.dtype != torch.float64:\n+        warnings.warn(\n+            \"Detected `nan` or `inf` value in computed metric, retrying computation in double precision\",\n+            UserWarning,\n+        )\n+        val = signal_distortion_ratio(preds.double(), target.double(), use_cg_iter, filter_length, zero_mean, load_diag)\n+\n     return val\n", "issue": "The sdr metric in TM sometime gives NaN for some input\n## \ud83d\udc1b Bug\r\n\r\nThis issue is related with `fast-bss-eval`'s torch version, see https://github.com/fakufaku/fast_bss_eval/issues/5\r\n\r\n### To Reproduce\r\n\r\n```\r\nimport numpy as np\r\nimport torch\r\n\r\nx = np.load('debug.npz')\r\npreds = torch.tensor(x['preds'])\r\ntarget = torch.tensor(x['target'])\r\nprint(preds.shape, target.shape)\r\n\r\nfrom torchmetrics.functional.audio import signal_distortion_ratio\r\nsdr = signal_distortion_ratio(preds, target)\r\nprint(sdr)\r\n\r\nfrom mir_eval.separation import bss_eval_sources\r\nsdr, _, _, _ = bss_eval_sources(target.numpy(), preds.numpy(), False)\r\nprint(sdr)\r\n```\r\n\r\noutputs:\r\n```\r\ntorch.Size([2, 64000]) torch.Size([2, 64000])\r\ntensor([-2.6815,     nan])\r\n[-2.68156071 44.58523729]\r\n```\r\n\r\nunzip [data.zip](https://github.com/PyTorchLightning/metrics/files/8309638/data.zip) to get the `debug.npz`\r\n\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n\r\n#### Code sample\r\n\r\n<!-- Ideally attach a minimal code sample to reproduce the decried issue.\r\nMinimal means having the shortest code but still preserving the bug. -->\r\n\r\n### Expected behavior\r\n\r\nthe results given by signal_distortion_ratio is close to the one given by `mir_eval`\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n### Environment\r\n\r\n- OS (e.g., Linux):\r\n- Python & PyTorch Version (e.g., 1.0):\r\n- How you installed PyTorch (`conda`, `pip`, build command if you used source):\r\n- Any other relevant information:\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\n\nfrom torchmetrics.utilities.imports import _FAST_BSS_EVAL_AVAILABLE, _TORCH_GREATER_EQUAL_1_8\n\nif _FAST_BSS_EVAL_AVAILABLE:\n    if _TORCH_GREATER_EQUAL_1_8:\n        from fast_bss_eval.torch.cgd import toeplitz_conjugate_gradient\n        from fast_bss_eval.torch.helpers import _normalize\n        from fast_bss_eval.torch.linalg import toeplitz\n        from fast_bss_eval.torch.metrics import compute_stats\n\n        solve = torch.linalg.solve\n    else:\n        import numpy\n        from fast_bss_eval.numpy.cgd import toeplitz_conjugate_gradient\n        from fast_bss_eval.numpy.helpers import _normalize\n        from fast_bss_eval.numpy.linalg import toeplitz\n        from fast_bss_eval.numpy.metrics import compute_stats\n\n        solve = numpy.linalg.solve\nelse:\n    toeplitz = None\n    toeplitz_conjugate_gradient = None\n    compute_stats = None\n    _normalize = None\n    __doctest_skip__ = [\"signal_distortion_ratio\"]\n\nfrom torch import Tensor\n\nfrom torchmetrics.utilities import rank_zero_warn\nfrom torchmetrics.utilities.checks import _check_same_shape\n\n\ndef signal_distortion_ratio(\n    preds: Tensor,\n    target: Tensor,\n    use_cg_iter: Optional[int] = None,\n    filter_length: int = 512,\n    zero_mean: bool = False,\n    load_diag: Optional[float] = None,\n) -> Tensor:\n    r\"\"\"Signal to Distortion Ratio (SDR) [1,2,3]\n\n    Args:\n        preds:\n            shape ``[..., time]``\n        target:\n            shape ``[..., time]``\n        use_cg_iter:\n            If provided, an iterative method is used to solve for the distortion\n            filter coefficients instead of direct Gaussian elimination.\n            This can speed up the computation of the metrics in case the filters\n            are long. Using a value of 10 here has been shown to provide\n            good accuracy in most cases and is sufficient when using this\n            loss to train neural separation networks.\n        filter_length:\n            The length of the distortion filter allowed\n        zero_mean:\n            When set to True, the mean of all signals is subtracted prior to computation of the metrics\n        load_diag:\n            If provided, this small value is added to the diagonal coefficients of\n            the system metrics when solving for the filter coefficients.\n            This can help stabilize the metric in the case where some of the reference\n            signals may sometimes be zero\n\n    Raises:\n        ModuleNotFoundError:\n            If ``fast-bss-eval`` package is not installed\n\n    Returns:\n        sdr value of shape ``[...]``\n\n    Example:\n\n        >>> from torchmetrics.functional.audio import signal_distortion_ratio\n        >>> import torch\n        >>> g = torch.manual_seed(1)\n        >>> preds = torch.randn(8000)\n        >>> target = torch.randn(8000)\n        >>> signal_distortion_ratio(preds, target)\n        tensor(-12.0589)\n        >>> # use with permutation_invariant_training\n        >>> from torchmetrics.functional.audio import permutation_invariant_training\n        >>> preds = torch.randn(4, 2, 8000)  # [batch, spk, time]\n        >>> target = torch.randn(4, 2, 8000)\n        >>> best_metric, best_perm = permutation_invariant_training(preds, target, signal_distortion_ratio, 'max')\n        >>> best_metric\n        tensor([-11.6375, -11.4358, -11.7148, -11.6325])\n        >>> best_perm\n        tensor([[1, 0],\n                [0, 1],\n                [1, 0],\n                [0, 1]])\n\n    .. note::\n       1. when pytorch<1.8.0, numpy will be used to calculate this metric, which causes ``sdr`` to be\n            non-differentiable and slower to calculate\n\n       2. using this metrics requires you to have ``fast-bss-eval`` install. Either install as ``pip install\n          torchmetrics[audio]`` or ``pip install fast-bss-eval``\n\n       3. preds and target need to have the same dtype, otherwise target will be converted to preds' dtype\n\n\n    References:\n        [1] Vincent, E., Gribonval, R., & Fevotte, C. (2006). Performance measurement in blind audio source separation.\n        IEEE Transactions on Audio, Speech and Language Processing, 14(4), 1462\u20131469.\n\n        [2] Scheibler, R. (2021). SDR -- Medium Rare with Fast Computations.\n\n        [3] https://github.com/fakufaku/fast_bss_eval\n    \"\"\"\n    if not _FAST_BSS_EVAL_AVAILABLE:\n        raise ModuleNotFoundError(\n            \"SDR metric requires that `fast-bss-eval` is installed.\"\n            \" Either install as `pip install torchmetrics[audio]` or `pip install fast-bss-eval`.\"\n        )\n    _check_same_shape(preds, target)\n\n    if not preds.dtype.is_floating_point:\n        preds = preds.float()  # for torch.norm\n\n    # half precision support\n    if preds.dtype == torch.float16:\n        preds = preds.to(torch.float32)\n\n    if preds.dtype != target.dtype:  # for torch.linalg.solve\n        target = target.to(preds.dtype)\n\n    if zero_mean:\n        preds = preds - preds.mean(dim=-1, keepdim=True)\n        target = target - target.mean(dim=-1, keepdim=True)\n\n    # normalize along time-axis\n    if not _TORCH_GREATER_EQUAL_1_8:\n        # use numpy if torch<1.8\n        rank_zero_warn(\n            \"Pytorch is under 1.8, thus SDR numpy version is used.\"\n            \"For better performance and differentiability, you should change to Pytorch v1.8 or above.\"\n        )\n        device = preds.device\n        preds = preds.detach().cpu().numpy()\n        target = target.detach().cpu().numpy()\n\n        preds = _normalize(preds, axis=-1)\n        target = _normalize(target, axis=-1)\n    else:\n        preds = _normalize(preds, dim=-1)\n        target = _normalize(target, dim=-1)\n\n    # solve for the optimal filter\n    # compute auto-correlation and cross-correlation\n    acf, xcorr = compute_stats(target, preds, length=filter_length, pairwise=False)\n\n    if load_diag is not None:\n        # the diagonal factor of the Toeplitz matrix is the first\n        # coefficient of the acf\n        acf[..., 0] += load_diag\n\n    if use_cg_iter is not None:\n        # use preconditioned conjugate gradient\n        sol = toeplitz_conjugate_gradient(acf, xcorr, n_iter=use_cg_iter)\n    else:\n        # regular matrix solver\n        r_mat = toeplitz(acf)\n        sol = solve(r_mat, xcorr)\n\n    # to tensor if torch<1.8\n    if not _TORCH_GREATER_EQUAL_1_8:\n        sol = torch.tensor(sol, device=device)\n        xcorr = torch.tensor(xcorr, device=device)\n\n    # compute the coherence\n    coh = torch.einsum(\"...l,...l->...\", xcorr, sol)\n\n    # transform to decibels\n    ratio = coh / (1 - coh)\n    val = 10.0 * torch.log10(ratio)\n    return val\n\n\ndef scale_invariant_signal_distortion_ratio(preds: Tensor, target: Tensor, zero_mean: bool = False) -> Tensor:\n    \"\"\"Calculates Scale-invariant signal-to-distortion ratio (SI-SDR) metric. The SI-SDR value is in general\n    considered an overall measure of how good a source sound.\n\n    Args:\n        preds:\n            shape ``[...,time]``\n        target:\n            shape ``[...,time]``\n        zero_mean:\n            If to zero mean target and preds or not\n\n    Returns:\n        si-sdr value of shape [...]\n\n    Example:\n        >>> from torchmetrics.functional.audio import scale_invariant_signal_distortion_ratio\n        >>> target = torch.tensor([3.0, -0.5, 2.0, 7.0])\n        >>> preds = torch.tensor([2.5, 0.0, 2.0, 8.0])\n        >>> scale_invariant_signal_distortion_ratio(preds, target)\n        tensor(18.4030)\n\n    References:\n        [1] Le Roux, Jonathan, et al. \"SDR half-baked or well done.\" IEEE International Conference on Acoustics, Speech\n        and Signal Processing (ICASSP) 2019.\n    \"\"\"\n    _check_same_shape(preds, target)\n    eps = torch.finfo(preds.dtype).eps\n\n    if zero_mean:\n        target = target - torch.mean(target, dim=-1, keepdim=True)\n        preds = preds - torch.mean(preds, dim=-1, keepdim=True)\n\n    alpha = (torch.sum(preds * target, dim=-1, keepdim=True) + eps) / (\n        torch.sum(target ** 2, dim=-1, keepdim=True) + eps\n    )\n    target_scaled = alpha * target\n\n    noise = target_scaled - preds\n\n    val = (torch.sum(target_scaled ** 2, dim=-1) + eps) / (torch.sum(noise ** 2, dim=-1) + eps)\n    val = 10 * torch.log10(val)\n\n    return val\n", "path": "torchmetrics/functional/audio/sdr.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport warnings\nfrom typing import Optional\n\nimport torch\n\nfrom torchmetrics.utilities.imports import _FAST_BSS_EVAL_AVAILABLE, _TORCH_GREATER_EQUAL_1_8\n\nif _FAST_BSS_EVAL_AVAILABLE:\n    if _TORCH_GREATER_EQUAL_1_8:\n        from fast_bss_eval.torch.cgd import toeplitz_conjugate_gradient\n        from fast_bss_eval.torch.helpers import _normalize\n        from fast_bss_eval.torch.linalg import toeplitz\n        from fast_bss_eval.torch.metrics import compute_stats\n\n        solve = torch.linalg.solve\n    else:\n        import numpy\n        from fast_bss_eval.numpy.cgd import toeplitz_conjugate_gradient\n        from fast_bss_eval.numpy.helpers import _normalize\n        from fast_bss_eval.numpy.linalg import toeplitz\n        from fast_bss_eval.numpy.metrics import compute_stats\n\n        solve = numpy.linalg.solve\nelse:\n    toeplitz = None\n    toeplitz_conjugate_gradient = None\n    compute_stats = None\n    _normalize = None\n    __doctest_skip__ = [\"signal_distortion_ratio\"]\n\nfrom torch import Tensor\n\nfrom torchmetrics.utilities import rank_zero_warn\nfrom torchmetrics.utilities.checks import _check_same_shape\n\n\ndef signal_distortion_ratio(\n    preds: Tensor,\n    target: Tensor,\n    use_cg_iter: Optional[int] = None,\n    filter_length: int = 512,\n    zero_mean: bool = False,\n    load_diag: Optional[float] = None,\n) -> Tensor:\n    r\"\"\"Signal to Distortion Ratio (SDR) [1,2,3]\n\n    Args:\n        preds:\n            shape ``[..., time]``\n        target:\n            shape ``[..., time]``\n        use_cg_iter:\n            If provided, an iterative method is used to solve for the distortion\n            filter coefficients instead of direct Gaussian elimination.\n            This can speed up the computation of the metrics in case the filters\n            are long. Using a value of 10 here has been shown to provide\n            good accuracy in most cases and is sufficient when using this\n            loss to train neural separation networks.\n        filter_length:\n            The length of the distortion filter allowed\n        zero_mean:\n            When set to True, the mean of all signals is subtracted prior to computation of the metrics\n        load_diag:\n            If provided, this small value is added to the diagonal coefficients of\n            the system metrics when solving for the filter coefficients.\n            This can help stabilize the metric in the case where some of the reference\n            signals may sometimes be zero\n\n    Raises:\n        ModuleNotFoundError:\n            If ``fast-bss-eval`` package is not installed\n\n    Returns:\n        sdr value of shape ``[...]``\n\n    Example:\n\n        >>> from torchmetrics.functional.audio import signal_distortion_ratio\n        >>> import torch\n        >>> g = torch.manual_seed(1)\n        >>> preds = torch.randn(8000)\n        >>> target = torch.randn(8000)\n        >>> signal_distortion_ratio(preds, target)\n        tensor(-12.0589)\n        >>> # use with permutation_invariant_training\n        >>> from torchmetrics.functional.audio import permutation_invariant_training\n        >>> preds = torch.randn(4, 2, 8000)  # [batch, spk, time]\n        >>> target = torch.randn(4, 2, 8000)\n        >>> best_metric, best_perm = permutation_invariant_training(preds, target, signal_distortion_ratio, 'max')\n        >>> best_metric\n        tensor([-11.6375, -11.4358, -11.7148, -11.6325])\n        >>> best_perm\n        tensor([[1, 0],\n                [0, 1],\n                [1, 0],\n                [0, 1]])\n\n    .. note::\n       1. when pytorch<1.8.0, numpy will be used to calculate this metric, which causes ``sdr`` to be\n            non-differentiable and slower to calculate\n\n       2. using this metrics requires you to have ``fast-bss-eval`` install. Either install as ``pip install\n          torchmetrics[audio]`` or ``pip install fast-bss-eval``\n\n       3. preds and target need to have the same dtype, otherwise target will be converted to preds' dtype\n\n\n    References:\n        [1] Vincent, E., Gribonval, R., & Fevotte, C. (2006). Performance measurement in blind audio source separation.\n        IEEE Transactions on Audio, Speech and Language Processing, 14(4), 1462\u20131469.\n\n        [2] Scheibler, R. (2021). SDR -- Medium Rare with Fast Computations.\n\n        [3] https://github.com/fakufaku/fast_bss_eval\n    \"\"\"\n    if not _FAST_BSS_EVAL_AVAILABLE:\n        raise ModuleNotFoundError(\n            \"SDR metric requires that `fast-bss-eval` is installed.\"\n            \" Either install as `pip install torchmetrics[audio]` or `pip install fast-bss-eval`.\"\n        )\n    _check_same_shape(preds, target)\n\n    if not preds.dtype.is_floating_point:\n        preds = preds.float()  # for torch.norm\n\n    # half precision support\n    if preds.dtype == torch.float16:\n        preds = preds.to(torch.float32)\n\n    if preds.dtype != target.dtype:  # for torch.linalg.solve\n        target = target.to(preds.dtype)\n\n    if zero_mean:\n        preds = preds - preds.mean(dim=-1, keepdim=True)\n        target = target - target.mean(dim=-1, keepdim=True)\n\n    # normalize along time-axis\n    if not _TORCH_GREATER_EQUAL_1_8:\n        # use numpy if torch<1.8\n        rank_zero_warn(\n            \"Pytorch is under 1.8, thus SDR numpy version is used.\"\n            \"For better performance and differentiability, you should change to Pytorch v1.8 or above.\"\n        )\n        device = preds.device\n        preds = preds.detach().cpu().numpy()\n        target = target.detach().cpu().numpy()\n\n        preds = _normalize(preds, axis=-1)\n        target = _normalize(target, axis=-1)\n    else:\n        preds = _normalize(preds, dim=-1)\n        target = _normalize(target, dim=-1)\n\n    # solve for the optimal filter\n    # compute auto-correlation and cross-correlation\n    acf, xcorr = compute_stats(target, preds, length=filter_length, pairwise=False)\n\n    if load_diag is not None:\n        # the diagonal factor of the Toeplitz matrix is the first\n        # coefficient of the acf\n        acf[..., 0] += load_diag\n\n    if use_cg_iter is not None:\n        # use preconditioned conjugate gradient\n        sol = toeplitz_conjugate_gradient(acf, xcorr, n_iter=use_cg_iter)\n    else:\n        # regular matrix solver\n        r_mat = toeplitz(acf)\n        sol = solve(r_mat, xcorr)\n\n    # to tensor if torch<1.8\n    if not _TORCH_GREATER_EQUAL_1_8:\n        sol = torch.tensor(sol, device=device)\n        xcorr = torch.tensor(xcorr, device=device)\n\n    # compute the coherence\n    coh = torch.einsum(\"...l,...l->...\", xcorr, sol)\n\n    # transform to decibels\n    ratio = coh / (1 - coh)\n    val = 10.0 * torch.log10(ratio)\n\n    # recompute sdr in float64 if val is NaN or Inf\n    if (torch.isnan(val).any() or torch.isinf(val).any()) and preds.dtype != torch.float64:\n        warnings.warn(\n            \"Detected `nan` or `inf` value in computed metric, retrying computation in double precision\",\n            UserWarning,\n        )\n        val = signal_distortion_ratio(preds.double(), target.double(), use_cg_iter, filter_length, zero_mean, load_diag)\n\n    return val\n\n\ndef scale_invariant_signal_distortion_ratio(preds: Tensor, target: Tensor, zero_mean: bool = False) -> Tensor:\n    \"\"\"Calculates Scale-invariant signal-to-distortion ratio (SI-SDR) metric. The SI-SDR value is in general\n    considered an overall measure of how good a source sound.\n\n    Args:\n        preds:\n            shape ``[...,time]``\n        target:\n            shape ``[...,time]``\n        zero_mean:\n            If to zero mean target and preds or not\n\n    Returns:\n        si-sdr value of shape [...]\n\n    Example:\n        >>> from torchmetrics.functional.audio import scale_invariant_signal_distortion_ratio\n        >>> target = torch.tensor([3.0, -0.5, 2.0, 7.0])\n        >>> preds = torch.tensor([2.5, 0.0, 2.0, 8.0])\n        >>> scale_invariant_signal_distortion_ratio(preds, target)\n        tensor(18.4030)\n\n    References:\n        [1] Le Roux, Jonathan, et al. \"SDR half-baked or well done.\" IEEE International Conference on Acoustics, Speech\n        and Signal Processing (ICASSP) 2019.\n    \"\"\"\n    _check_same_shape(preds, target)\n    eps = torch.finfo(preds.dtype).eps\n\n    if zero_mean:\n        target = target - torch.mean(target, dim=-1, keepdim=True)\n        preds = preds - torch.mean(preds, dim=-1, keepdim=True)\n\n    alpha = (torch.sum(preds * target, dim=-1, keepdim=True) + eps) / (\n        torch.sum(target ** 2, dim=-1, keepdim=True) + eps\n    )\n    target_scaled = alpha * target\n\n    noise = target_scaled - preds\n\n    val = (torch.sum(target_scaled ** 2, dim=-1) + eps) / (torch.sum(noise ** 2, dim=-1) + eps)\n    val = 10 * torch.log10(val)\n\n    return val\n", "path": "torchmetrics/functional/audio/sdr.py"}]}