docstr: PITLoss Wrapper function. Similar to espnet2/enh/loss/wrapper/pit_solver.py
  
  Args:
      inf: Iterable[torch.Tensor], (batch, num_inf, ...)
      inf_lens: Iterable[torch.Tensor], (batch, num_inf, ...)
      ref: Iterable[torch.Tensor], (batch, num_ref, ...)
      ref_lens: Iterable[torch.Tensor], (batch, num_ref, ...)
      permute_inf: If true, permute the inference and inference_lens according to
          the optimal permutation.
docstr_tok: ["PITLoss","Wrapper","function",".","Similar","to","espnet2\/enh\/loss\/wrapper\/pit_solver.py","Args",":","inf",":","Iterable","[","torch.Tensor","]",",","(","batch",",","num_inf",",","...",")","inf_lens",":","Iterable","[","torch.Tensor","]",",","(","batch",",","num_inf",",","...",")","ref",":","Iterable","[","torch.Tensor","]",",","(","batch",",","num_ref",",","...",")","ref_lens",":","Iterable","[","torch.Tensor","]",",","(","batch",",","num_ref",",","...",")","permute_inf",":","If","true",",","permute","the","inference","and","inference_lens","according","to","the","optimal","permutation","."]
code: def forward(
      self,
      inf: torch.Tensor,
      inf_lens: torch.Tensor,
      ref: torch.Tensor,
      ref_lens: torch.Tensor,
      others: Dict = None,
  ):
      """PITLoss Wrapper function. Similar to espnet2/enh/loss/wrapper/pit_solver.py
  
      Args:
          inf: Iterable[torch.Tensor], (batch, num_inf, ...)
          inf_lens: Iterable[torch.Tensor], (batch, num_inf, ...)
          ref: Iterable[torch.Tensor], (batch, num_ref, ...)
          ref_lens: Iterable[torch.Tensor], (batch, num_ref, ...)
          permute_inf: If true, permute the inference and inference_lens according to
              the optimal permutation.
      """
      assert (
          self.num_ref
          == inf.shape[1]
          == inf_lens.shape[1]
          == ref.shape[1]
          == ref_lens.shape[1]
      ), (
          self.num_ref,
          inf.shape,
          inf_lens.shape,
          ref.shape,
          ref_lens.shape,
      )
  
      all_permutations = torch.as_tensor(
          list(
              itertools.permutations(
                  range(self.num_ref), r=self.num_ref
              )
          )
      )
  
      stats = defaultdict(list)
  
      def pre_hook(func, *args, **kwargs):
          ret = func(*args, **kwargs)
          for k, v in getattr(self.criterion_fn, "stats", {}).items():
              stats[k].append(v)
          return ret
  
      def pair_loss(permutation):
          return sum(
              [
                  pre_hook(
                      self.criterion_fn,
                      inf[:, j],
                      inf_lens[:, j],
                      ref[:, i],
                      ref_lens[:, i],
                  )
                  for i, j in enumerate(permutation)
              ]
          ) / len(permutation)
  
      losses = torch.stack(
          [pair_loss(p) for p in all_permutations], dim=1
      )  # (batch_size, num_perm)
  
      min_losses, min_ids = torch.min(losses, dim=1)
      min_ids = (
          min_ids.cpu()
      )  # because all_permutations is a cpu tensor.
      opt_perm = all_permutations[min_ids]  # (batch_size, num_ref)
  
      # Permute the inf and inf_lens according to the optimal perm
      return min_losses.mean(), opt_perm
  
imports: import itertools
  from collections import defaultdict
  from contextlib import contextmanager
  from typing import Callable, Dict, List, Optional, Tuple, Union
  import torch
  from packaging.version import parse
  from typeguard import check_argument_types
  from espnet2.asr.ctc import CTC
  from espnet2.asr.decoder.abs_decoder import AbsDecoder
  from espnet2.asr.encoder.abs_encoder import AbsEncoder
  from espnet2.asr.espnet_model import ESPnetASRModel
  from espnet2.asr.frontend.abs_frontend import AbsFrontend
  from espnet2.asr.postencoder.abs_postencoder import AbsPostEncoder
  from espnet2.asr.preencoder.abs_preencoder import AbsPreEncoder
  from espnet2.asr.specaug.abs_specaug import AbsSpecAug
  from espnet2.enh.loss.wrappers.abs_wrapper import AbsLossWrapper
  from espnet2.layers.abs_normalize import AbsNormalize
  from espnet2.torch_utils.device_funcs import force_gatherable
label_desc: Use image node_id 2 for calling the PITLossWrapper obj's underlying member method code with example usage: obj.forward(inf, inf_lens, ref, ref_lens, others) and returns: min_losses,  opt_perm, ret, unknown