code: def make_segmentation_mask(
      size=DEFAULT_SIZE,
      *,
      num_categories=10,
      batch_dims=(),
      dtype=None,
      device="cpu",
  ):
      """Make a "segmentation" mask, i.e. (*, H, W), where the category is encoded as pixel value"""
      return tv_tensors.Mask(
          torch.testing.make_tensor(
              (*batch_dims, *size),
              low=0,
              high=num_categories,
              dtype=dtype or torch.uint8,
              device=device,
          )
      )
  
imports: import contextlib
  import functools
  import itertools
  import os
  import pathlib
  import random
  import re
  import shutil
  import sys
  import tempfile
  import warnings
  from subprocess import CalledProcessError, check_output, STDOUT
  import numpy
  import PIL.Image
  import pytest
  import torch
  import torch.testing
  from PIL import Image
  from torch.testing._comparison import BooleanPair, NonePair, not_close_error_metas, NumberPair, TensorLikePair
  from torchvision import io, tv_tensors
  from torchvision.transforms._functional_tensor import _max_value
  from torchvision.transforms.v2.functional import to_image, to_pil_image
label_desc: Use image node_id 28 for calling a global function with example usage: make_segmentation_mask(size) and returns: tv_tensors