code: def gen_test_batch(
      batch_size: int,
      label_name: str = "label",
      prediction_name: str = "prediction",
      weight_name: str = "weight",
      tensor_name: str = "tensor",
      mask_tensor_name: Optional[str] = None,
      label_value: Optional[torch.Tensor] = None,
      prediction_value: Optional[torch.Tensor] = None,
      weight_value: Optional[torch.Tensor] = None,
      mask: Optional[torch.Tensor] = None,
      n_classes: Optional[int] = None,
  ) -> Dict[str, torch.Tensor]:
      if label_value is not None:
          label = label_value
      else:
          label = torch.randint(
              0, n_classes or 2, (batch_size,)
          ).double()
      if prediction_value is not None:
          prediction = prediction_value
      else:
          prediction = (
              torch.rand(batch_size, dtype=torch.double)
              if n_classes is None
              else torch.rand(batch_size, n_classes, dtype=torch.double)
          )
      if weight_value is not None:
          weight = weight_value
      else:
          weight = torch.rand(batch_size, dtype=torch.double)
      test_batch = {
          label_name: label,
          prediction_name: prediction,
          weight_name: weight,
          tensor_name: torch.rand(batch_size, dtype=torch.double),
      }
      if mask_tensor_name is not None:
          if mask is None:
              mask = torch.ones(batch_size, dtype=torch.double)
          test_batch[mask_tensor_name] = mask
  
      return test_batch
  
imports: import abc
  import os
  import random
  import tempfile
  import uuid
  from typing import Any, Callable, Dict, List, Optional, Tuple, Type
  from unittest.mock import Mock, patch
  import torch
  import torch.distributed
  import torch.distributed.launcher
  from torchrec.metrics.auc import AUCMetric
  from torchrec.metrics.auprc import AUPRCMetric
  from torchrec.metrics.model_utils import parse_task_model_outputs
  from torchrec.metrics.rec_metric import RecComputeMode, RecMetric, RecTaskInfo
label_desc: Use image node_id 1 for calling a global function with example usage: gen_test_batch(batch_size, label_name, prediction_name, weight_name, tensor_name, mask_tensor_name, label_value, prediction_value, weight_value, mask, n_classes) and returns: test_batch