We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
pit bug: the device and type of preds and target may not the same with the results returned by metric_func
## 🐛 Bug

the device and type of preds and target may not the same with the results returned by metric_func

### To Reproduce

use a self-designed metric-func which accepts complex tensor but return real tensor

<!-- 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

<!-- A clear and concise description of what you expected to happen. -->

### Environment

- PyTorch Version (e.g., 1.0):
- OS (e.g., Linux):
- How you installed PyTorch (`conda`, `pip`, source):
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version:
- GPU models and configuration:
- 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/pit.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 import warnings
15 from itertools import permutations
16 from typing import Any, Callable, Dict, Tuple, Union
17 
18 import torch
19 from torch import Tensor
20 
21 from torchmetrics.utilities.checks import _check_same_shape
22 from torchmetrics.utilities.imports import _SCIPY_AVAILABLE
23 
24 # _ps_dict: cache of permutations
25 # it's necessary to cache it, otherwise it will consume a large amount of time
26 _ps_dict: dict = {}  # _ps_dict[str(spk_num)+str(device)] = permutations
27 
28 
29 def _find_best_perm_by_linear_sum_assignment(
30     metric_mtx: torch.Tensor,
31     eval_func: Union[torch.min, torch.max],
32 ) -> Tuple[Tensor, Tensor]:
33     """Solves the linear sum assignment problem using scipy, and returns the best metric values and the
34     corresponding permutations.
35 
36     Args:
37         metric_mtx:
38             the metric matrix, shape [batch_size, spk_num, spk_num]
39         eval_func:
40             the function to reduce the metric values of different the permutations
41 
42     Returns:
43         best_metric:
44             shape [batch]
45         best_perm:
46             shape [batch, spk]
47     """
48     from scipy.optimize import linear_sum_assignment
49 
50     mmtx = metric_mtx.detach().cpu()
51     best_perm = torch.tensor([linear_sum_assignment(pwm, eval_func == torch.max)[1] for pwm in mmtx])
52     best_perm = best_perm.to(metric_mtx.device)
53     best_metric = torch.gather(metric_mtx, 2, best_perm[:, :, None]).mean([-1, -2])
54     return best_metric, best_perm  # shape [batch], shape [batch, spk]
55 
56 
57 def _find_best_perm_by_exhuastive_method(
58     metric_mtx: torch.Tensor,
59     eval_func: Union[torch.min, torch.max],
60 ) -> Tuple[Tensor, Tensor]:
61     """Solves the linear sum assignment problem using exhuastive method, i.e. exhuastively calculates the metric
62     values of all possible permutations, and returns the best metric values and the corresponding permutations.
63 
64     Args:
65         metric_mtx:
66             the metric matrix, shape [batch_size, spk_num, spk_num]
67         eval_func:
68             the function to reduce the metric values of different the permutations
69 
70     Returns:
71         best_metric:
72             shape [batch]
73         best_perm:
74             shape [batch, spk]
75     """
76     # create/read/cache the permutations and its indexes
77     # reading from cache would be much faster than creating in CPU then moving to GPU
78     batch_size, spk_num = metric_mtx.shape[:2]
79     key = str(spk_num) + str(metric_mtx.device)
80     if key not in _ps_dict:
81         # ps: all the permutations, shape [spk_num, perm_num]
82         # ps: In i-th permutation, the predcition corresponds to the j-th target is ps[j,i]
83         ps = torch.tensor(list(permutations(range(spk_num))), device=metric_mtx.device).T
84         _ps_dict[key] = ps
85     else:
86         ps = _ps_dict[key]  # all the permutations, shape [spk_num, perm_num]
87 
88     # find the metric of each permutation
89     perm_num = ps.shape[-1]
90     # shape [batch_size, spk_num, perm_num]
91     bps = ps[None, ...].expand(batch_size, spk_num, perm_num)
92     # shape [batch_size, spk_num, perm_num]
93     metric_of_ps_details = torch.gather(metric_mtx, 2, bps)
94     # shape [batch_size, perm_num]
95     metric_of_ps = metric_of_ps_details.mean(dim=1)
96 
97     # find the best metric and best permutation
98     best_metric, best_indexes = eval_func(metric_of_ps, dim=1)
99     best_indexes = best_indexes.detach()
100     best_perm = ps.T[best_indexes, :]
101     return best_metric, best_perm  # shape [batch], shape [batch, spk]
102 
103 
104 def pit(
105     preds: torch.Tensor, target: torch.Tensor, metric_func: Callable, eval_func: str = "max", **kwargs: Dict[str, Any]
106 ) -> Tuple[Tensor, Tensor]:
107     """Permutation invariant training (PIT). The PIT implements the famous Permutation Invariant Training method.
108 
109     [1] in speech separation field in order to calculate audio metrics in a permutation invariant way.
110 
111     Args:
112         preds:
113             shape [batch, spk, ...]
114         target:
115             shape [batch, spk, ...]
116         metric_func:
117             a metric function accept a batch of target and estimate,
118             i.e. metric_func(preds[:, i, ...], target[:, j, ...]), and returns a batch of metric tensors [batch]
119         eval_func:
120             the function to find the best permutation, can be 'min' or 'max',
121             i.e. the smaller the better or the larger the better.
122         kwargs:
123             additional args for metric_func
124 
125     Returns:
126         best_metric of shape [batch],
127         best_perm of shape [batch]
128 
129     Example:
130         >>> from torchmetrics.functional.audio import si_sdr
131         >>> # [batch, spk, time]
132         >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])
133         >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])
134         >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')
135         >>> best_metric
136         tensor([-5.1091])
137         >>> best_perm
138         tensor([[0, 1]])
139         >>> pit_permutate(preds, best_perm)
140         tensor([[[-0.0579,  0.3560, -0.9604],
141                  [-0.1719,  0.3205,  0.2951]]])
142 
143     Reference:
144         [1]	`Permutation Invariant Training of Deep Models`_
145     """
146     _check_same_shape(preds, target)
147     if eval_func not in ["max", "min"]:
148         raise ValueError(f'eval_func can only be "max" or "min" but got {eval_func}')
149     if target.ndim < 2:
150         raise ValueError(f"Inputs must be of shape [batch, spk, ...], got {target.shape} and {preds.shape} instead")
151 
152     # calculate the metric matrix
153     batch_size, spk_num = target.shape[0:2]
154     metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=preds.dtype, device=target.device)
155     for t in range(spk_num):
156         for e in range(spk_num):
157             metric_mtx[:, t, e] = metric_func(preds[:, e, ...], target[:, t, ...], **kwargs)
158 
159     # find best
160     op = torch.max if eval_func == "max" else torch.min
161     if spk_num < 3 or not _SCIPY_AVAILABLE:
162         if spk_num >= 3 and not _SCIPY_AVAILABLE:
163             warnings.warn(
164                 f"In pit metric for speaker-num {spk_num}>3, we recommend installing scipy for better performance"
165             )
166 
167         best_metric, best_perm = _find_best_perm_by_exhuastive_method(metric_mtx, op)
168     else:
169         best_metric, best_perm = _find_best_perm_by_linear_sum_assignment(metric_mtx, op)
170 
171     return best_metric, best_perm
172 
173 
174 def pit_permutate(preds: Tensor, perm: Tensor) -> Tensor:
175     """permutate estimate according to perm.
176 
177     Args:
178         preds (Tensor): the estimates you want to permutate, shape [batch, spk, ...]
179         perm (Tensor): the permutation returned from pit, shape [batch, spk]
180 
181     Returns:
182         Tensor: the permutated version of estimate
183 
184     Example:
185         >>> from torchmetrics.functional.audio import si_sdr
186         >>> # [batch, spk, time]
187         >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])
188         >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])
189         >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')
190         >>> best_metric
191         tensor([-5.1091])
192         >>> best_perm
193         tensor([[0, 1]])
194         >>> pit_permutate(preds, best_perm)
195         tensor([[[-0.0579,  0.3560, -0.9604],
196                  [-0.1719,  0.3205,  0.2951]]])
197     """
198     preds_pmted = torch.stack([torch.index_select(pred, 0, p) for pred, p in zip(preds, perm)])
199     return preds_pmted
200
```
--- 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/pit.py b/torchmetrics/functional/audio/pit.py
--- a/torchmetrics/functional/audio/pit.py
+++ b/torchmetrics/functional/audio/pit.py
@@ -151,10 +151,17 @@
 
     # calculate the metric matrix
     batch_size, spk_num = target.shape[0:2]
-    metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=preds.dtype, device=target.device)
-    for t in range(spk_num):
-        for e in range(spk_num):
-            metric_mtx[:, t, e] = metric_func(preds[:, e, ...], target[:, t, ...], **kwargs)
+    metric_mtx = None
+    for target_idx in range(spk_num):  # we have spk_num speeches in target in each sample
+        for preds_idx in range(spk_num):  # we have spk_num speeches in preds in each sample
+            if metric_mtx is not None:
+                metric_mtx[:, target_idx, preds_idx] = metric_func(
+                    preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs
+                )
+            else:
+                first_ele = metric_func(preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs)
+                metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=first_ele.dtype, device=first_ele.device)
+                metric_mtx[:, target_idx, preds_idx] = first_ele
 
     # find best
     op = torch.max if eval_func == "max" else torch.min

{"golden_diff": "diff --git a/torchmetrics/functional/audio/pit.py b/torchmetrics/functional/audio/pit.py\n--- a/torchmetrics/functional/audio/pit.py\n+++ b/torchmetrics/functional/audio/pit.py\n@@ -151,10 +151,17 @@\n \n     # calculate the metric matrix\n     batch_size, spk_num = target.shape[0:2]\n-    metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=preds.dtype, device=target.device)\n-    for t in range(spk_num):\n-        for e in range(spk_num):\n-            metric_mtx[:, t, e] = metric_func(preds[:, e, ...], target[:, t, ...], **kwargs)\n+    metric_mtx = None\n+    for target_idx in range(spk_num):  # we have spk_num speeches in target in each sample\n+        for preds_idx in range(spk_num):  # we have spk_num speeches in preds in each sample\n+            if metric_mtx is not None:\n+                metric_mtx[:, target_idx, preds_idx] = metric_func(\n+                    preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs\n+                )\n+            else:\n+                first_ele = metric_func(preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs)\n+                metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=first_ele.dtype, device=first_ele.device)\n+                metric_mtx[:, target_idx, preds_idx] = first_ele\n \n     # find best\n     op = torch.max if eval_func == \"max\" else torch.min\n", "issue": "pit bug: the device and type of preds and target may not the same with the results returned by metric_func\n## \ud83d\udc1b Bug\r\n\r\nthe device and type of preds and target may not the same with the results returned by metric_func\r\n\r\n### To Reproduce\r\n\r\nuse a self-designed metric-func which accepts complex tensor but return real tensor\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\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n### Environment\r\n\r\n- PyTorch Version (e.g., 1.0):\r\n- OS (e.g., Linux):\r\n- How you installed PyTorch (`conda`, `pip`, source):\r\n- Build command you used (if compiling from source):\r\n- Python version:\r\n- CUDA/cuDNN version:\r\n- GPU models and configuration:\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.\nimport warnings\nfrom itertools import permutations\nfrom typing import Any, Callable, Dict, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.checks import _check_same_shape\nfrom torchmetrics.utilities.imports import _SCIPY_AVAILABLE\n\n# _ps_dict: cache of permutations\n# it's necessary to cache it, otherwise it will consume a large amount of time\n_ps_dict: dict = {}  # _ps_dict[str(spk_num)+str(device)] = permutations\n\n\ndef _find_best_perm_by_linear_sum_assignment(\n    metric_mtx: torch.Tensor,\n    eval_func: Union[torch.min, torch.max],\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Solves the linear sum assignment problem using scipy, and returns the best metric values and the\n    corresponding permutations.\n\n    Args:\n        metric_mtx:\n            the metric matrix, shape [batch_size, spk_num, spk_num]\n        eval_func:\n            the function to reduce the metric values of different the permutations\n\n    Returns:\n        best_metric:\n            shape [batch]\n        best_perm:\n            shape [batch, spk]\n    \"\"\"\n    from scipy.optimize import linear_sum_assignment\n\n    mmtx = metric_mtx.detach().cpu()\n    best_perm = torch.tensor([linear_sum_assignment(pwm, eval_func == torch.max)[1] for pwm in mmtx])\n    best_perm = best_perm.to(metric_mtx.device)\n    best_metric = torch.gather(metric_mtx, 2, best_perm[:, :, None]).mean([-1, -2])\n    return best_metric, best_perm  # shape [batch], shape [batch, spk]\n\n\ndef _find_best_perm_by_exhuastive_method(\n    metric_mtx: torch.Tensor,\n    eval_func: Union[torch.min, torch.max],\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Solves the linear sum assignment problem using exhuastive method, i.e. exhuastively calculates the metric\n    values of all possible permutations, and returns the best metric values and the corresponding permutations.\n\n    Args:\n        metric_mtx:\n            the metric matrix, shape [batch_size, spk_num, spk_num]\n        eval_func:\n            the function to reduce the metric values of different the permutations\n\n    Returns:\n        best_metric:\n            shape [batch]\n        best_perm:\n            shape [batch, spk]\n    \"\"\"\n    # create/read/cache the permutations and its indexes\n    # reading from cache would be much faster than creating in CPU then moving to GPU\n    batch_size, spk_num = metric_mtx.shape[:2]\n    key = str(spk_num) + str(metric_mtx.device)\n    if key not in _ps_dict:\n        # ps: all the permutations, shape [spk_num, perm_num]\n        # ps: In i-th permutation, the predcition corresponds to the j-th target is ps[j,i]\n        ps = torch.tensor(list(permutations(range(spk_num))), device=metric_mtx.device).T\n        _ps_dict[key] = ps\n    else:\n        ps = _ps_dict[key]  # all the permutations, shape [spk_num, perm_num]\n\n    # find the metric of each permutation\n    perm_num = ps.shape[-1]\n    # shape [batch_size, spk_num, perm_num]\n    bps = ps[None, ...].expand(batch_size, spk_num, perm_num)\n    # shape [batch_size, spk_num, perm_num]\n    metric_of_ps_details = torch.gather(metric_mtx, 2, bps)\n    # shape [batch_size, perm_num]\n    metric_of_ps = metric_of_ps_details.mean(dim=1)\n\n    # find the best metric and best permutation\n    best_metric, best_indexes = eval_func(metric_of_ps, dim=1)\n    best_indexes = best_indexes.detach()\n    best_perm = ps.T[best_indexes, :]\n    return best_metric, best_perm  # shape [batch], shape [batch, spk]\n\n\ndef pit(\n    preds: torch.Tensor, target: torch.Tensor, metric_func: Callable, eval_func: str = \"max\", **kwargs: Dict[str, Any]\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Permutation invariant training (PIT). The PIT implements the famous Permutation Invariant Training method.\n\n    [1] in speech separation field in order to calculate audio metrics in a permutation invariant way.\n\n    Args:\n        preds:\n            shape [batch, spk, ...]\n        target:\n            shape [batch, spk, ...]\n        metric_func:\n            a metric function accept a batch of target and estimate,\n            i.e. metric_func(preds[:, i, ...], target[:, j, ...]), and returns a batch of metric tensors [batch]\n        eval_func:\n            the function to find the best permutation, can be 'min' or 'max',\n            i.e. the smaller the better or the larger the better.\n        kwargs:\n            additional args for metric_func\n\n    Returns:\n        best_metric of shape [batch],\n        best_perm of shape [batch]\n\n    Example:\n        >>> from torchmetrics.functional.audio import si_sdr\n        >>> # [batch, spk, time]\n        >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])\n        >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])\n        >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')\n        >>> best_metric\n        tensor([-5.1091])\n        >>> best_perm\n        tensor([[0, 1]])\n        >>> pit_permutate(preds, best_perm)\n        tensor([[[-0.0579,  0.3560, -0.9604],\n                 [-0.1719,  0.3205,  0.2951]]])\n\n    Reference:\n        [1]\t`Permutation Invariant Training of Deep Models`_\n    \"\"\"\n    _check_same_shape(preds, target)\n    if eval_func not in [\"max\", \"min\"]:\n        raise ValueError(f'eval_func can only be \"max\" or \"min\" but got {eval_func}')\n    if target.ndim < 2:\n        raise ValueError(f\"Inputs must be of shape [batch, spk, ...], got {target.shape} and {preds.shape} instead\")\n\n    # calculate the metric matrix\n    batch_size, spk_num = target.shape[0:2]\n    metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=preds.dtype, device=target.device)\n    for t in range(spk_num):\n        for e in range(spk_num):\n            metric_mtx[:, t, e] = metric_func(preds[:, e, ...], target[:, t, ...], **kwargs)\n\n    # find best\n    op = torch.max if eval_func == \"max\" else torch.min\n    if spk_num < 3 or not _SCIPY_AVAILABLE:\n        if spk_num >= 3 and not _SCIPY_AVAILABLE:\n            warnings.warn(\n                f\"In pit metric for speaker-num {spk_num}>3, we recommend installing scipy for better performance\"\n            )\n\n        best_metric, best_perm = _find_best_perm_by_exhuastive_method(metric_mtx, op)\n    else:\n        best_metric, best_perm = _find_best_perm_by_linear_sum_assignment(metric_mtx, op)\n\n    return best_metric, best_perm\n\n\ndef pit_permutate(preds: Tensor, perm: Tensor) -> Tensor:\n    \"\"\"permutate estimate according to perm.\n\n    Args:\n        preds (Tensor): the estimates you want to permutate, shape [batch, spk, ...]\n        perm (Tensor): the permutation returned from pit, shape [batch, spk]\n\n    Returns:\n        Tensor: the permutated version of estimate\n\n    Example:\n        >>> from torchmetrics.functional.audio import si_sdr\n        >>> # [batch, spk, time]\n        >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])\n        >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])\n        >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')\n        >>> best_metric\n        tensor([-5.1091])\n        >>> best_perm\n        tensor([[0, 1]])\n        >>> pit_permutate(preds, best_perm)\n        tensor([[[-0.0579,  0.3560, -0.9604],\n                 [-0.1719,  0.3205,  0.2951]]])\n    \"\"\"\n    preds_pmted = torch.stack([torch.index_select(pred, 0, p) for pred, p in zip(preds, perm)])\n    return preds_pmted\n", "path": "torchmetrics/functional/audio/pit.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.\nimport warnings\nfrom itertools import permutations\nfrom typing import Any, Callable, Dict, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.checks import _check_same_shape\nfrom torchmetrics.utilities.imports import _SCIPY_AVAILABLE\n\n# _ps_dict: cache of permutations\n# it's necessary to cache it, otherwise it will consume a large amount of time\n_ps_dict: dict = {}  # _ps_dict[str(spk_num)+str(device)] = permutations\n\n\ndef _find_best_perm_by_linear_sum_assignment(\n    metric_mtx: torch.Tensor,\n    eval_func: Union[torch.min, torch.max],\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Solves the linear sum assignment problem using scipy, and returns the best metric values and the\n    corresponding permutations.\n\n    Args:\n        metric_mtx:\n            the metric matrix, shape [batch_size, spk_num, spk_num]\n        eval_func:\n            the function to reduce the metric values of different the permutations\n\n    Returns:\n        best_metric:\n            shape [batch]\n        best_perm:\n            shape [batch, spk]\n    \"\"\"\n    from scipy.optimize import linear_sum_assignment\n\n    mmtx = metric_mtx.detach().cpu()\n    best_perm = torch.tensor([linear_sum_assignment(pwm, eval_func == torch.max)[1] for pwm in mmtx])\n    best_perm = best_perm.to(metric_mtx.device)\n    best_metric = torch.gather(metric_mtx, 2, best_perm[:, :, None]).mean([-1, -2])\n    return best_metric, best_perm  # shape [batch], shape [batch, spk]\n\n\ndef _find_best_perm_by_exhuastive_method(\n    metric_mtx: torch.Tensor,\n    eval_func: Union[torch.min, torch.max],\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Solves the linear sum assignment problem using exhuastive method, i.e. exhuastively calculates the metric\n    values of all possible permutations, and returns the best metric values and the corresponding permutations.\n\n    Args:\n        metric_mtx:\n            the metric matrix, shape [batch_size, spk_num, spk_num]\n        eval_func:\n            the function to reduce the metric values of different the permutations\n\n    Returns:\n        best_metric:\n            shape [batch]\n        best_perm:\n            shape [batch, spk]\n    \"\"\"\n    # create/read/cache the permutations and its indexes\n    # reading from cache would be much faster than creating in CPU then moving to GPU\n    batch_size, spk_num = metric_mtx.shape[:2]\n    key = str(spk_num) + str(metric_mtx.device)\n    if key not in _ps_dict:\n        # ps: all the permutations, shape [spk_num, perm_num]\n        # ps: In i-th permutation, the predcition corresponds to the j-th target is ps[j,i]\n        ps = torch.tensor(list(permutations(range(spk_num))), device=metric_mtx.device).T\n        _ps_dict[key] = ps\n    else:\n        ps = _ps_dict[key]  # all the permutations, shape [spk_num, perm_num]\n\n    # find the metric of each permutation\n    perm_num = ps.shape[-1]\n    # shape [batch_size, spk_num, perm_num]\n    bps = ps[None, ...].expand(batch_size, spk_num, perm_num)\n    # shape [batch_size, spk_num, perm_num]\n    metric_of_ps_details = torch.gather(metric_mtx, 2, bps)\n    # shape [batch_size, perm_num]\n    metric_of_ps = metric_of_ps_details.mean(dim=1)\n\n    # find the best metric and best permutation\n    best_metric, best_indexes = eval_func(metric_of_ps, dim=1)\n    best_indexes = best_indexes.detach()\n    best_perm = ps.T[best_indexes, :]\n    return best_metric, best_perm  # shape [batch], shape [batch, spk]\n\n\ndef pit(\n    preds: torch.Tensor, target: torch.Tensor, metric_func: Callable, eval_func: str = \"max\", **kwargs: Dict[str, Any]\n) -> Tuple[Tensor, Tensor]:\n    \"\"\"Permutation invariant training (PIT). The PIT implements the famous Permutation Invariant Training method.\n\n    [1] in speech separation field in order to calculate audio metrics in a permutation invariant way.\n\n    Args:\n        preds:\n            shape [batch, spk, ...]\n        target:\n            shape [batch, spk, ...]\n        metric_func:\n            a metric function accept a batch of target and estimate,\n            i.e. metric_func(preds[:, i, ...], target[:, j, ...]), and returns a batch of metric tensors [batch]\n        eval_func:\n            the function to find the best permutation, can be 'min' or 'max',\n            i.e. the smaller the better or the larger the better.\n        kwargs:\n            additional args for metric_func\n\n    Returns:\n        best_metric of shape [batch],\n        best_perm of shape [batch]\n\n    Example:\n        >>> from torchmetrics.functional.audio import si_sdr\n        >>> # [batch, spk, time]\n        >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])\n        >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])\n        >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')\n        >>> best_metric\n        tensor([-5.1091])\n        >>> best_perm\n        tensor([[0, 1]])\n        >>> pit_permutate(preds, best_perm)\n        tensor([[[-0.0579,  0.3560, -0.9604],\n                 [-0.1719,  0.3205,  0.2951]]])\n\n    Reference:\n        [1]\t`Permutation Invariant Training of Deep Models`_\n    \"\"\"\n    _check_same_shape(preds, target)\n    if eval_func not in [\"max\", \"min\"]:\n        raise ValueError(f'eval_func can only be \"max\" or \"min\" but got {eval_func}')\n    if target.ndim < 2:\n        raise ValueError(f\"Inputs must be of shape [batch, spk, ...], got {target.shape} and {preds.shape} instead\")\n\n    # calculate the metric matrix\n    batch_size, spk_num = target.shape[0:2]\n    metric_mtx = None\n    for target_idx in range(spk_num):  # we have spk_num speeches in target in each sample\n        for preds_idx in range(spk_num):  # we have spk_num speeches in preds in each sample\n            if metric_mtx is not None:\n                metric_mtx[:, target_idx, preds_idx] = metric_func(\n                    preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs\n                )\n            else:\n                first_ele = metric_func(preds[:, preds_idx, ...], target[:, target_idx, ...], **kwargs)\n                metric_mtx = torch.empty((batch_size, spk_num, spk_num), dtype=first_ele.dtype, device=first_ele.device)\n                metric_mtx[:, target_idx, preds_idx] = first_ele\n\n    # find best\n    op = torch.max if eval_func == \"max\" else torch.min\n    if spk_num < 3 or not _SCIPY_AVAILABLE:\n        if spk_num >= 3 and not _SCIPY_AVAILABLE:\n            warnings.warn(\n                f\"In pit metric for speaker-num {spk_num}>3, we recommend installing scipy for better performance\"\n            )\n\n        best_metric, best_perm = _find_best_perm_by_exhuastive_method(metric_mtx, op)\n    else:\n        best_metric, best_perm = _find_best_perm_by_linear_sum_assignment(metric_mtx, op)\n\n    return best_metric, best_perm\n\n\ndef pit_permutate(preds: Tensor, perm: Tensor) -> Tensor:\n    \"\"\"permutate estimate according to perm.\n\n    Args:\n        preds (Tensor): the estimates you want to permutate, shape [batch, spk, ...]\n        perm (Tensor): the permutation returned from pit, shape [batch, spk]\n\n    Returns:\n        Tensor: the permutated version of estimate\n\n    Example:\n        >>> from torchmetrics.functional.audio import si_sdr\n        >>> # [batch, spk, time]\n        >>> preds = torch.tensor([[[-0.0579,  0.3560, -0.9604], [-0.1719,  0.3205,  0.2951]]])\n        >>> target = torch.tensor([[[ 1.0958, -0.1648,  0.5228], [-0.4100,  1.1942, -0.5103]]])\n        >>> best_metric, best_perm = pit(preds, target, si_sdr, 'max')\n        >>> best_metric\n        tensor([-5.1091])\n        >>> best_perm\n        tensor([[0, 1]])\n        >>> pit_permutate(preds, best_perm)\n        tensor([[[-0.0579,  0.3560, -0.9604],\n                 [-0.1719,  0.3205,  0.2951]]])\n    \"\"\"\n    preds_pmted = torch.stack([torch.index_select(pred, 0, p) for pred, p in zip(preds, perm)])\n    return preds_pmted\n", "path": "torchmetrics/functional/audio/pit.py"}]}