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","]",")","."]
docstr: Args:
      batch (Dict[str, torch.Tensor]): currently expects input dense features
          to be mapped to the key "float_features" and input sparse features
          to be mapped to the key "id_list_features".
  
  Returns:
      Dict[str, torch.Tensor]: output of inference.
docstr_tok: ["Args",":","batch","(","Dict","[","str",",","torch.Tensor","]",")",":","currently","expects","input","dense","features","to","be","mapped","to","the","key","``","float_features","''","and","input","sparse","features","to","be","mapped","to","the","key","``","id_list_features","''",".","Returns",":","Dict","[","str",",","torch.Tensor","]",":","output","of","inference","."]
code: def predict_forward(
      self, batch: Dict[str, torch.Tensor]
  ) -> Dict[str, torch.Tensor]:
      """
      Args:
          batch (Dict[str, torch.Tensor]): currently expects input dense features
              to be mapped to the key "float_features" and input sparse features
              to be mapped to the key "id_list_features".
  
      Returns:
          Dict[str, torch.Tensor]: output of inference.
      """
  
      try:
          logits = self.predict_module(
              batch["float_features"],
              KeyedJaggedTensor(
                  keys=self.id_list_features_keys,
                  lengths=batch["id_list_features.lengths"],
                  values=batch["id_list_features.values"],
              ),
          )
          predictions = logits.sigmoid()
      except Exception as e:
          logger.info(e)
          raise e
  
      # Flip predictions tensor to be 1D. TODO: Determine why prediction shape
      # can be 2D at times (likely due to input format?)
      predictions = predictions.reshape(
          [
              predictions.size()[0],
          ]
      )
  
      return {
          "default": predictions.to(
              torch.device("cpu"), non_blocking=True
          ).float()
      }
  
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 2 for calling the DLRMPredictModule obj's underlying member method code with example usage: obj.predict_forward(batch) and returns: dict