class_docstr: Simple TwoTower (UV) Model. Embeds two different entities into the same space.
  A simplified version of the `A Dual Augmented Two-tower Model for Online Large-scale Recommendation
  <https://dlp-kdd.github.io/assets/pdf/DLP-KDD_2021_paper_4.pdf>`_ model.
  Used to train the retrieval model
  
  Embeddings trained with this model will be indexed and queried in the retrieval example.
  
  Args:
      embedding_bag_collection (EmbeddingBagCollection): embedding_bag_collection with two EmbeddingBags
      layer_sizes (List[int]): list of the layer_sizes for the MLP
      device (Optional[torch.device])
  
  Example::
  
      m = TwoTower(ebc, [16, 8], device)
docstr: Args:
      kjt (KeyedJaggedTensor): KJT containing query_ids and candidate_ids to query
  
  Returns:
      Tuple[torch.Tensor, torch.Tensor]: Tuple containing embeddings for each tower
docstr_tok: ["Args",":","kjt","(","KeyedJaggedTensor",")",":","KJT","containing","query_ids","and","candidate_ids","to","query","Returns",":","Tuple","[","torch.Tensor",",","torch.Tensor","]",":","Tuple","containing","embeddings","for","each","tower"]
code: def forward(
      self, kjt: KeyedJaggedTensor
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Args:
          kjt (KeyedJaggedTensor): KJT containing query_ids and candidate_ids to query
  
      Returns:
          Tuple[torch.Tensor, torch.Tensor]: Tuple containing embeddings for each tower
      """
      pooled_embeddings = self.ebc(kjt)
      query_embedding: torch.Tensor = self.query_proj(
          torch.cat(
              [
                  pooled_embeddings[feature]
                  for feature in self._feature_names_query
              ],
              dim=1,
          )
      )
      candidate_embedding: torch.Tensor = self.candidate_proj(
          torch.cat(
              [
                  pooled_embeddings[feature]
                  for feature in self._candidate_feature_names
              ],
              dim=1,
          )
      )
      return query_embedding, candidate_embedding
  
imports: from typing import Any, List, Mapping, Optional, OrderedDict, Tuple, Union
  import faiss
  import faiss.contrib.torch_utils
  import torch
  from torch import nn
  from torchrec.datasets.utils import Batch
  from torchrec.modules.embedding_modules import EmbeddingBagCollection
  from torchrec.modules.mlp import MLP
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 2 for calling the TwoTower obj's underlying member method code with example usage: obj.forward(kjt) and returns: query_embedding,  candidate_embedding