code: def test_pitwrapper(device):
      from speechbrain.nnet.losses import PitWrapper
      import torch
      from torch import nn
  
      base_loss = nn.MSELoss(reduction="none")
      pit = PitWrapper(base_loss)
      predictions = torch.rand(
          (2, 32, 4), device=device
      )  # batch, frames, sources
      p = (3, 0, 2, 1)
      # same but we invert the ordering to check if permutation invariant
      targets = predictions[..., p]
      loss, opt_p = pit(predictions, targets)
      assert [x == p for x in opt_p] == [
          True for i in range(len(opt_p))
      ]
      predictions = pit.reorder_tensor(predictions, opt_p)
      assert torch.all(torch.eq(base_loss(predictions, targets), 0))
  
      predictions = torch.rand(
          (3, 32, 32, 32, 5), device=device
      )  # batch, ..., sources
      p = (3, 0, 2, 1, 4)
      targets = predictions[
          ..., p
      ]  # same but we invert the ordering to check if permutation invariant
      loss, opt_p = pit(predictions, targets)
      assert [x == p for x in opt_p] == [
          True for i in range(len(opt_p))
      ]
      predictions = pit.reorder_tensor(predictions, opt_p)
      assert torch.all(torch.eq(base_loss(predictions, targets), 0))
  
file_path: speechbrain/tests/unittests/test_losses.py
imports: import torch
  import pytest
label_desc: Use image node_id 6 for calling a global function with example usage: test_pitwrapper(device) without return types