code: def test_multi_batch_image_gradients(
      batch_size=5, height=5, width=5, channels=1
  ):
      """Test whether the module correctly calculates gradients for known input with non-unity batch size."""
      single_channel_img = torch.arange(
          0, 1 * height * width * channels, dtype=torch.float32
      )
      single_channel_img = torch.reshape(
          single_channel_img, (channels, height, width)
      )
      image = torch.stack(
          [single_channel_img for _ in range(batch_size)], dim=0
      )
  
      true_dy = [
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [5.0, 5.0, 5.0, 5.0, 5.0],
          [0.0, 0.0, 0.0, 0.0, 0.0],
      ]
      true_dy = Tensor(true_dy)
  
      dy, dx = image_gradients(image)
  
      for batch_id in range(batch_size):
          assert torch.allclose(dy[batch_id, 0, :, :], true_dy)
      assert dy.shape == (batch_size, 1, height, width)
      assert dx.shape == (batch_size, 1, height, width)
  
imports: import pytest
  import torch
  from torch import Tensor
  from torchmetrics.functional import image_gradients
label_desc: Use image node_id 3 for calling a global function with example usage: test_multi_batch_image_gradients(batch_size, height, width, channels) without return types