code: def init_model(
      eb_configs: List[EmbeddingBagConfig],
      dense_arch_layer_sizes: Optional[List[int]],
      over_arch_layer_sizes: Optional[List[int]],
      learning_rate: float,
      batch_size: int,
      device: torch.device,
  ) -> torch.nn.Module:
      dlrm_model = DLRM(
          embedding_bag_collection=EmbeddingBagCollection(
              tables=eb_configs, device=torch.device("meta")
          ),
          dense_in_features=len(DEFAULT_INT_NAMES),
          dense_arch_layer_sizes=dense_arch_layer_sizes or [],
          over_arch_layer_sizes=over_arch_layer_sizes or [],
          dense_device=device,
      )
  
      train_model = DLRMTrain(dlrm_model)
  
      _apply_optimizer_in_backward(
          torch.optim.SGD,
          train_model.model.sparse_arch.parameters(),
          {"lr": learning_rate},
      )
  
      planner = EmbeddingShardingPlanner(
          topology=Topology(
              local_world_size=get_local_size(),
              world_size=dist.get_world_size(),
              compute_device=device.type,
          ),
          batch_size=batch_size,
          storage_reservation=HeuristicalStorageReservation(
              percentage=0.05
          ),
      )
      plan = planner.collective_plan(
          train_model, get_default_sharders(), dist.GroupMember.WORLD
      )
  
      model = DistributedModelParallel(
          module=train_model,
          device=device,
          plan=plan,
      )
  
      return model
  
imports: import argparse
  import sys
  import tempfile
  import uuid
  from typing import Iterator, List, Optional, Tuple
  import torch
  from pyre_extensions import none_throws
  from torch import distributed
  from torch.distributed import launcher
  from torch.distributed.optim.apply_optimizer_in_backward import _apply_optimizer_in_backward
  from torch.utils.data import DataLoader
  from torcheval.metrics.classification.auroc import BinaryAUROC
  from torcheval.metrics.toolkit import sync_and_compute
  from torchrec import EmbeddingBagCollection
  from torchrec.datasets.criteo import DEFAULT_CAT_NAMES, DEFAULT_INT_NAMES
  from torchrec.datasets.random import RandomRecDataset
  from torchrec.distributed import TrainPipelineSparseDist
  from torchrec.distributed.comm import get_local_size
  from torchrec.distributed.model_parallel import DistributedModelParallel, get_default_sharders
  from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
  from torchrec.distributed.planner.storage_reservations import HeuristicalStorageReservation
  from torchrec.models.dlrm import DLRM, DLRMTrain
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.optim.keyed import KeyedOptimizerWrapper
  from torchrec.optim.optimizers import in_backward_optimizer_filter
  from torchtnt.framework.callbacks import TQDMProgressBar
  from torchtnt.framework.fit import fit
  from torchtnt.framework.state import State
  from torchtnt.framework.unit import EvalUnit, TrainUnit
  from torchtnt.utils import get_process_group_backend_from_device, init_from_env, rank_zero_print
  from torchtnt.utils.loggers.tensorboard import TensorBoardLogger
label_desc: Use image node_id 4 for calling a global function with example usage: init_model(eb_configs, dense_arch_layer_sizes, over_arch_layer_sizes, learning_rate, batch_size, device) and returns: model