code: def _test_pipelining(
      self,
      sharder: EmbeddingBagCollectionSharder,
      execute_all_batches: bool,
  ) -> None:
      pipeline = self._setup_pipeline(sharder, execute_all_batches)
      cpu_model, cpu_optimizer = self._setup_cpu_model_and_opt()
      data = self._generate_data()
  
      dataloader = iter(data)
      if not execute_all_batches:
          data = data[:-2]
  
      for batch in data:
          cpu_optimizer.zero_grad()
          loss, pred = cpu_model(batch)
          loss.backward()
          cpu_optimizer.step()
  
          pred_gpu = pipeline.progress(dataloader)
  
          self.assertEqual(len(pipeline._pipelined_modules), 2)
          self.assertEqual(pred_gpu.device, self.device)
          self.assertEqual(pred_gpu.cpu().size(), pred.size())
  
imports: import copy
  import os
  import unittest
  from dataclasses import dataclass
  from typing import cast, Dict, List, Optional, Tuple
  import torch
  import torch.distributed
  from hypothesis import given, settings, strategies, Verbosity
  from torch import nn, optim
  from torch.distributed.fsdp import FullyShardedDataParallel
  from torchrec.distributed import DistributedModelParallel
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, KJTList
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionContext, EmbeddingBagCollectionSharder, ShardedEmbeddingBagCollection
  from torchrec.distributed.fp_embeddingbag import FeatureProcessedEmbeddingBagCollectionSharder, ShardedFeatureProcessedEmbeddingBagCollection
  from torchrec.distributed.sharding_plan import construct_module_sharding_plan, table_wise
  from torchrec.distributed.test_utils.test_model import ModelInput, TestEBCSharder, TestSparseNN
  from torchrec.distributed.test_utils.test_sharding import copy_state_dict
  from torchrec.distributed.train_pipeline import TrainPipelineBase, TrainPipelineSparseDist
  from torchrec.distributed.types import Awaitable, ModuleSharder, ParameterSharding, ShardingEnv, ShardingPlan, ShardingType
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.optim.keyed import KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
  from torchrec.streamable import Pipelineable
  from torchrec.test_utils import get_free_port, init_distributed_single_host
  from .test_fp_embeddingbag_utils import create_module_and_freeze
label_desc: Use image node_id 8 for calling the TrainPipelineSparseDistTest obj's underlying member method code with example usage: obj._test_pipelining(sharder, execute_all_batches) without return types