code: def test_flow_to_image(batch):
      h, w = 100, 100
      flow = torch.meshgrid(
          torch.arange(h), torch.arange(w), indexing="ij"
      )
      flow = torch.stack(flow[::-1], dim=0).float()
      flow[0] -= h / 2
      flow[1] -= w / 2
  
      if batch:
          flow = torch.stack([flow, flow])
  
      img = utils.flow_to_image(flow)
      assert img.shape == (2, 3, h, w) if batch else (3, h, w)
  
      path = os.path.join(
          os.path.dirname(os.path.abspath(__file__)),
          "assets",
          "expected_flow.pt",
      )
      expected_img = torch.load(
          path, map_location="cpu", weights_only=True
      )
  
      if batch:
          expected_img = torch.stack([expected_img, expected_img])
  
      assert_equal(expected_img, img)
  
imports: import os
  import re
  import sys
  import tempfile
  from io import BytesIO
  import numpy
  import pytest
  import torch
  import torchvision.transforms.functional
  import torchvision.utils
  from common_utils import assert_equal, cpu_and_cuda
  from PIL import __version__, Image, ImageColor
  from torchvision.transforms.v2.functional import to_dtype
label_desc: Use image node_id 21 for calling a global function with example usage: test_flow_to_image(batch) without return types