code: def test_semantic_segmentation_heads_registry(head):
      img = torch.rand(1, 3, 32, 32)
      backbone = SEMANTIC_SEGMENTATION_BACKBONES.get("resnet50")(
          pretrained=False
      )
      head = SEMANTIC_SEGMENTATION_HEADS.get(head)(
          backbone=backbone, num_classes=10
      )
      assert backbone
      assert head
      head.eval()
      res = head(img)
      if isinstance(res, dict):
          res = res["out"]
      assert res.shape[1] == 10
  
file_path: lightning-flash/tests/image/semantic_segm/test_heads.py
imports: import unittest
  import pytest
  import torch
  from flash.core.utilities.imports import _SEGMENTATION_MODELS_AVAILABLE
  from flash.image.segmentation import SemanticSegmentation
  from flash.image.segmentation.backbones import SEMANTIC_SEGMENTATION_BACKBONES
  from flash.image.segmentation.heads import SEMANTIC_SEGMENTATION_HEADS
label_desc: Use image node_id 1 for calling a global function with example usage: test_semantic_segmentation_heads_registry(head) without return types