code: def make_sample_test_batch(
      raw_data_path, processed_data_path, device
  ):
      """Create the test_data_loader and sample a batch from it. This batch will be used
      to measure the forward pass of the model throughout this experiment.
      """
      test_data_loader = make_test_data_loader(
          raw_data_path, processed_data_path
      )
  
      test_iter = iter(test_data_loader)
  
      test_batch = next(test_iter)
  
      X_test, lS_o_test, lS_i_test, _, _, _ = unpack_batch(test_batch)
  
      X, lS_o, lS_i = dlrm_wrap(X_test, lS_o_test, lS_i_test, device)
      batch = {"X": X, "lS_o": lS_o, "lS_i": lS_i}
  
      return batch
  
file_path: pytorch/torch/ao/pruning/_experimental/data_sparsifier/benchmarks/evaluate_forward_time.py
imports: from typing import Dict, List
  import torch
  from dlrm_s_pytorch import unpack_batch
  import numpy
  import time
  from dlrm_utils import make_test_data_loader, fetch_model, dlrm_wrap
  import pandas
  import argparse
label_desc: Use image node_id 2 for calling a global function with example usage: make_sample_test_batch(raw_data_path, processed_data_path, device) and returns: batch