code: def test_inference_universal_segmentation_head(self):
      model = (
          OneFormerForUniversalSegmentation.from_pretrained(
              self.model_checkpoints
          )
          .to(torch_device)
          .eval()
      )
      processor = self.default_processor
      image = prepare_img()
      inputs = processor(image, ["semantic"], return_tensors="pt").to(
          torch_device
      )
      inputs_shape = inputs["pixel_values"].shape
      # check size
      self.assertEqual(inputs_shape, (1, 3, 512, 682))
  
      with torch.no_grad():
          outputs = model(**inputs)
  
      # masks_queries_logits
      masks_queries_logits = outputs.masks_queries_logits
      self.assertEqual(
          masks_queries_logits.shape,
          (
              1,
              model.config.num_queries,
              inputs_shape[-2] // 4,
              (inputs_shape[-1] + 2) // 4,
          ),
      )
      expected_slice = [
          [
              [3.1848, 4.2141, 4.1993],
              [2.9000, 3.5721, 3.6603],
              [2.5358, 3.0883, 3.6168],
          ]
      ]
      expected_slice = torch.tensor(expected_slice).to(torch_device)
      self.assertTrue(
          torch.allclose(
              masks_queries_logits[0, 0, :3, :3],
              expected_slice,
              atol=TOLERANCE,
          )
      )
      # class_queries_logits
      class_queries_logits = outputs.class_queries_logits
      self.assertEqual(
          class_queries_logits.shape,
          (1, model.config.num_queries, model.config.num_labels + 1),
      )
      expected_slice = torch.tensor(
          [
              [3.0668, -1.1833, -5.1103],
              [3.3440, -3.3620, -5.1101],
              [2.6017, -4.3613, -4.1444],
          ]
      ).to(torch_device)
      self.assertTrue(
          torch.allclose(
              class_queries_logits[0, :3, :3],
              expected_slice,
              atol=TOLERANCE,
          )
      )
  
imports: import copy
  import inspect
  import unittest
  import numpy
  from tests.test_modeling_common import floats_tensor
  from transformers import OneFormerConfig, is_torch_available, is_vision_available
  from transformers.testing_utils import require_torch, require_torch_accelerator, require_torch_fp16, require_torch_multi_gpu, require_vision, slow, torch_device
  from transformers.utils import cached_property
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin
  from ...test_pipeline_mixin import PipelineTesterMixin
label_desc: Use image node_id 4 for calling the OneFormerModelIntegrationTest obj's underlying member method code with example usage: obj.test_inference_universal_segmentation_head() without return types