code: def _test_sharded_forward(
      self,
      world_size: int,
      model_class: TestSparseNNBase,
      embedding_groups: Dict[str, List[str]],
      tables: List[EmbeddingTableConfig],
      sharders: List[ModuleSharder[nn.Module]],
      quantize_callable: Callable[[nn.Module], nn.Module],
      quantize_callable_kwargs: Dict[str, Any],
      dedup_features_names: Optional[List[str]] = None,
      dedup_tables: Optional[List[EmbeddingTableConfig]] = None,
      weighted_tables: Optional[List[EmbeddingTableConfig]] = None,
      constraints: Optional[Dict[str, ParameterConstraints]] = None,
      generate: ModelInputCallable = ModelInput.generate,
  ) -> None:
      default_rank = 0
      cuda_device = torch.device(f"cuda:{default_rank}")
      torch.cuda.set_device(cuda_device)
  
      # Generate model & inputs.
      (global_model, _inputs) = gen_model_and_input(
          model_class=model_class,
          tables=tables,
          weighted_tables=weighted_tables,
          dedup_tables=dedup_tables,
          dedup_feature_names=dedup_features_names,
          embedding_groups=embedding_groups,
          world_size=1,  # generate only one copy of feature for inference
          num_float_features=16,
          dense_device=cuda_device,
          sparse_device=cuda_device,
          generate=generate,
          long_indices=False,
      )
      global_model = quantize_callable(
          global_model, **quantize_callable_kwargs
      )
      local_input = _inputs[0][1][default_rank].to(cuda_device)
  
      # Shard model.
      if dedup_features_names:
          local_model = model_class(
              tables=cast(
                  List[BaseEmbeddingConfig],
                  tables + dedup_tables if dedup_tables else tables,
              ),
              weighted_tables=cast(
                  List[BaseEmbeddingConfig], weighted_tables
              ),
              dedup_feature_names=dedup_features_names,
              embedding_groups=embedding_groups,
              dense_device=cuda_device,
              sparse_device=torch.device("meta"),
              num_float_features=16,
          )
      else:
          local_model = model_class(
              tables=cast(
                  List[BaseEmbeddingConfig],
                  tables,
              ),
              weighted_tables=cast(
                  List[BaseEmbeddingConfig], weighted_tables
              ),
              embedding_groups=embedding_groups,
              dense_device=cuda_device,
              sparse_device=torch.device("meta"),
              num_float_features=16,
          )
      local_model = quantize_callable(
          local_model, **quantize_callable_kwargs
      )
  
      planner = EmbeddingShardingPlanner(
          topology=Topology(world_size, "cuda"),
          constraints=constraints,
      )
      plan: ShardingPlan = planner.plan(local_model, sharders)
  
      # Generate a sharded model on a default rank.
      local_model = DistributedModelParallel(
          module=local_model,
          env=ShardingEnv.from_local(world_size, default_rank),
          plan=plan,
          sharders=sharders,
          init_data_parallel=False,
      )
  
      # materialize inference sharded model on one device for dense part
      local_model = local_model.copy(cuda_device)
  
      copy_state_dict(
          local_model.state_dict(), global_model.state_dict()
      )
  
      # Run a single training step of the sharded model.
      with torch.inference_mode():
          shard_pred = local_model(local_input)
  
      # Run second training step of the unsharded model.
      with torch.inference_mode():
          global_pred = global_model(local_input)
  
      # Compare predictions of sharded vs unsharded models.
      rtol, atol = _get_default_rtol_and_atol(global_pred, shard_pred)
      torch.testing.assert_close(
          global_pred, shard_pred, rtol=rtol, atol=atol
      )
  
imports: import os
  import unittest
  from collections import defaultdict
  from typing import Any, Callable, cast, Dict, List, Optional, OrderedDict, Tuple
  import numpy
  import torch
  import torch.nn
  from fbgemm_gpu.split_embedding_configs import EmbOptimType
  from hypothesis import given, settings, strategies, Verbosity
  from torch import distributed
  from torch.distributed._shard.sharded_tensor import ShardedTensor
  from torchrec import distributed
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel, EmbeddingTableConfig
  from torchrec.distributed.embeddingbag import EmbeddingBagCollectionSharder, ShardedEmbeddingBagCollection
  from torchrec.distributed.fused_embeddingbag import ShardedFusedEmbeddingBagCollection
  from torchrec.distributed.model_parallel import DistributedModelParallel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, ParameterConstraints, Topology
  from torchrec.distributed.sharding_plan import get_default_sharders
  from torchrec.distributed.test_utils.test_model import _get_default_rtol_and_atol, ModelInput, TestSparseNN, TestSparseNNBase
  from torchrec.distributed.test_utils.test_sharding import copy_state_dict, create_test_sharder, gen_model_and_input, ModelInputCallable, SharderType
  from torchrec.distributed.types import ModuleSharder, ShardingEnv, ShardingPlan, ShardingType
  from torchrec.modules.embedding_configs import BaseEmbeddingConfig, EmbeddingBagConfig, PoolingType
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.fused_embedding_modules import FusedEmbeddingBagCollection
  from torchrec.test_utils import get_free_port, seed_and_log
label_desc: Use image node_id 3 for calling the InferenceModelParallelTestBase obj's underlying member method code with example usage: obj._test_sharded_forward(world_size, model_class, embedding_groups, tables, sharders, quantize_callable, quantize_callable_kwargs, dedup_features_names, dedup_tables, weighted_tables, constraints, generate) without return types