code: def test_unet_discriminator(
      shape, norm_type, output_channels, downsampling_factor
  ):
      image = unet_image(shape)
      model = UnetDiscriminator(
          output_downsampling_factor=downsampling_factor,
          output_channels=output_channels,
          norm_type=norm_type,
      )
  
      expected_shape = (
          shape[0],
          output_channels,
          shape[2] // downsampling_factor,
          shape[3] // downsampling_factor,
      )
  
      if norm_type == "identity":
          test_model = model.model[0]
  
          while not isinstance(test_model, nn.Identity):
              assert isinstance(test_model.left_block[0], nn.Identity)
              assert isinstance(test_model.left_block[3], nn.Identity)
  
              test_model = test_model.child
      if norm_type == "spectral":
          test_model = model.model[0]
  
          while not isinstance(test_model, nn.Identity):
              assert nn.utils.parametrize.is_parametrized(
                  test_model.left_block[2]
              )
              assert nn.utils.parametrize.is_parametrized(
                  test_model.left_block[5]
              )
  
              test_model = test_model.child
      if norm_type == "instance":
          test_model = model.model[0]
          top = True
  
          while not isinstance(test_model, nn.Identity):
              if top:
                  top = False
              else:
                  assert isinstance(
                      test_model.left_block[0], nn.InstanceNorm2d
                  )
              assert isinstance(
                  test_model.left_block[3], nn.InstanceNorm2d
              )
  
              test_model = test_model.child
  
      output = model(image)
  
      assert output.shape == expected_shape
  
imports: import pytest
  import torch
  import torch.nn
  from neuralcompression.models import ConditionalUnetDiscriminator, UnetDiscriminator
label_desc: Use image node_id 2 for calling a global function with example usage: test_unet_discriminator(shape, norm_type, output_channels, downsampling_factor) without return types