class_docstr: nn.Module to wrap DLRM model to use for inference.
  
  Args:
      embedding_bag_collection (EmbeddingBagCollection): collection of embedding bags
          used to define SparseArch.
      dense_in_features (int): the dimensionality of the dense input features.
      dense_arch_layer_sizes (List[int]): the layer sizes for the DenseArch.
      over_arch_layer_sizes (List[int]): the layer sizes for the OverArch. NOTE: The
          output dimension of the InteractionArch should not be manually specified
          here.
      id_list_features_keys (List[str]): the names of the sparse features. Used to
          construct a batch for inference.
      dense_device: (Optional[torch.device]).
class_docstr_tok: ["nn.Module","to","wrap","DLRM","model","to","use","for","inference",".","Args",":","embedding_bag_collection","(","EmbeddingBagCollection",")",":","collection","of","embedding","bags","used","to","define","SparseArch",".","dense_in_features","(","int",")",":","the","dimensionality","of","the","dense","input","features",".","dense_arch_layer_sizes","(","List","[","int","]",")",":","the","layer","sizes","for","the","DenseArch",".","over_arch_layer_sizes","(","List","[","int","]",")",":","the","layer","sizes","for","the","OverArch",".","NOTE",":","The","output","dimension","of","the","InteractionArch","should","not","be","manually","specified","here",".","id_list_features_keys","(","List","[","str","]",")",":","the","names","of","the","sparse","features",".","Used","to","construct","a","batch","for","inference",".","dense_device",":","(","Optional","[","torch.device","]",")","."]
code: def __init__(
      self,
      embedding_bag_collection: EmbeddingBagCollection,
      dense_in_features: int,
      dense_arch_layer_sizes: List[int],
      over_arch_layer_sizes: List[int],
      id_list_features_keys: List[str],
      dense_device: Optional[torch.device] = None,
  ) -> None:
      module = DLRM(
          embedding_bag_collection=embedding_bag_collection,
          dense_in_features=dense_in_features,
          dense_arch_layer_sizes=dense_arch_layer_sizes,
          over_arch_layer_sizes=over_arch_layer_sizes,
          dense_device=dense_device,
      )
      super().__init__(module)
  
      self.id_list_features_keys: List[str] = id_list_features_keys
  
imports: import logging
  from dataclasses import dataclass
  from typing import Dict, List, Optional
  import torch
  import torchrec.distributed
  from torchrec.distributed.embedding_types import EmbeddingComputeKernel
  from torchrec.distributed.planner import EmbeddingShardingPlanner, ParameterConstraints, Topology
  from torchrec.distributed.quant_embeddingbag import QuantEmbeddingBagCollectionSharder
  from torchrec.distributed.types import ShardingType
  from torchrec.inference.model_packager import load_pickle_config
  from torchrec.inference.modules import PredictFactory, PredictModule, quantize_embeddings
  from torchrec.models.dlrm import DLRM
  from torchrec.modules.embedding_configs import EmbeddingBagConfig
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 1 to create a new DLRMPredictModule object from inherited base classes: PredictModule with example: obj = DLRMPredictModule(embedding_bag_collection, dense_in_features, dense_arch_layer_sizes, over_arch_layer_sizes, id_list_features_keys, dense_device)