class_docstr: nn.Module to wrap DLRM model to use with train_pipeline.
  
  DLRM Recsys model from "Deep Learning Recommendation Model for Personalization and
  Recommendation Systems" (https://arxiv.org/abs/1906.00091). Processes sparse
  features by learning pooled embeddings for each feature. Learns the relationship
  between dense features and sparse features by projecting dense features into the
  same embedding space. Also, learns the pairwise relationships between sparse
  features.
  
  The module assumes all sparse features have the same embedding dimension
  (i.e, each EmbeddingBagConfig uses the same embedding_dim)
  
  Args:
      dlrm_module: DLRM module (DLRM or DLRM_Projection or DLRM_DCN) to be used in
      training
  
  Example::
  
      ebc = EmbeddingBagCollection(config=ebc_config)
      dlrm_module = DLRM(
         embedding_bag_collection=ebc,
         dense_in_features=100,
         dense_arch_layer_sizes=[20],
         over_arch_layer_sizes=[5, 1],
      )
      dlrm_model = DLRMTrain(dlrm_module)
class_docstr_tok: ["nn.Module","to","wrap","DLRM","model","to","use","with","train_pipeline",".","DLRM","Recsys","model","from","``","Deep","Learning","Recommendation","Model","for","Personalization","and","Recommendation","Systems","''","(","https",":","\/\/arxiv.org\/abs\/1906.00091",")",".","Processes","sparse","features","by","learning","pooled","embeddings","for","each","feature",".","Learns","the","relationship","between","dense","features","and","sparse","features","by","projecting","dense","features","into","the","same","embedding","space",".","Also",",","learns","the","pairwise","relationships","between","sparse","features",".","The","module","assumes","all","sparse","features","have","the","same","embedding","dimension","(","i.e",",","each","EmbeddingBagConfig","uses","the","same","embedding_dim",")","Args",":","dlrm_module",":","DLRM","module","(","DLRM","or","DLRM_Projection","or","DLRM_DCN",")","to","be","used","in","training","Example",":",":","ebc","=","EmbeddingBagCollection","(","config=ebc_config",")","dlrm_module","=","DLRM","(","embedding_bag_collection=ebc",",","dense_in_features=100",",","dense_arch_layer_sizes=","[","20","]",",","over_arch_layer_sizes=","[","5",",","1","]",",",")","dlrm_model","=","DLRMTrain","(","dlrm_module",")"]
docstr: Args:
      batch: batch used with criteo and random data from torchrec.datasets
  Returns:
      Tuple[loss, Tuple[loss, logits, labels]]
docstr_tok: ["Args",":","batch",":","batch","used","with","criteo","and","random","data","from","torchrec.datasets","Returns",":","Tuple","[","loss",",","Tuple","[","loss",",","logits",",","labels","]","]"]
code: def forward(
      self, batch: Batch
  ) -> Tuple[
      torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
  ]:
      """
      Args:
          batch: batch used with criteo and random data from torchrec.datasets
      Returns:
          Tuple[loss, Tuple[loss, logits, labels]]
      """
      logits = self.model(batch.dense_features, batch.sparse_features)
      logits = logits.squeeze(-1)
      loss = self.loss_fn(logits, batch.labels.float())
  
      return loss, (
          loss.detach(),
          logits.detach(),
          batch.labels.detach(),
      )
  
imports: from typing import Dict, List, Optional, Tuple
  import torch
  from torch import nn
  from torchrec.datasets.utils import Batch
  from torchrec.modules.crossnet import LowRankCrossNet
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.mlp import MLP
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor, KeyedTensor
label_desc: Use image node_id 2 for calling the DLRMTrain obj's underlying member method code with example usage: obj.forward(batch) and returns: loss,  loss,  logits,  batch