class_docstr: Bidirectional Encoder = Transformer (self-attention)
  Transformer = MultiHead_Attention + Feed_Forward with sublayer connection
  
  Args:
      hidden (int): hidden size of transformer
      attn_heads (int): head sizes of multi-head attention
      feed_forward_hidden (int): feed_forward_hidden, usually 4*hidden_size
      dropout (float): the dropout probability
      device: (Optional[torch.device]).
  
  Example::
  
      self.transformer_blocks = nn.ModuleList(
          [
              TransformerBlock(emb_dim, nhead, emb_dim * 4, dropout, device=device)
              for _ in range(num_layers)
          ]
      )
class_docstr_tok: ["Bidirectional","Encoder","=","Transformer","(","self-attention",")","Transformer","=","MultiHead_Attention","+","Feed_Forward","with","sublayer","connection","Args",":","hidden","(","int",")",":","hidden","size","of","transformer","attn_heads","(","int",")",":","head","sizes","of","multi-head","attention","feed_forward_hidden","(","int",")",":","feed_forward_hidden",",","usually","4","*","hidden_size","dropout","(","float",")",":","the","dropout","probability","device",":","(","Optional","[","torch.device","]",")",".","Example",":",":","self.transformer_blocks","=","nn.ModuleList","(","[","TransformerBlock","(","emb_dim",",","nhead",",","emb_dim","*","4",",","dropout",",","device=device",")","for","_","in","range","(","num_layers",")","]",")"]
code: def __init__(
      self,
      hidden: int,
      attn_heads: int,
      feed_forward_hidden: int,
      dropout: float,
      device: Optional[torch.device] = None,
  ) -> None:
      super().__init__()
      self.attention = MultiHeadedAttention(
          num_heads=attn_heads,
          dim_model=hidden,
          dropout=dropout,
          device=device,
      )
      self.feed_forward = PositionwiseFeedForward(
          dim_model=hidden,
          d_ff=feed_forward_hidden,
          dropout=dropout,
          device=device,
      )
      self.input_sublayer = SublayerConnection(
          size=hidden, dropout=dropout, device=device
      )
      self.output_sublayer = SublayerConnection(
          size=hidden, dropout=dropout, device=device
      )
      self.dropout = nn.Dropout(p=dropout)
  
imports: import copy
  import math
  from typing import Callable, Optional, Tuple
  import torch
  import torch.nn
  from torchrec.modules.embedding_configs import EmbeddingConfig
  from torchrec.modules.embedding_modules import EmbeddingCollection
  from torchrec.sparse.jagged_tensor import KeyedJaggedTensor
label_desc: Use image node_id 1 to create a new TransformerBlock object from inherited base classes: nn with example: obj = TransformerBlock(hidden, attn_heads, feed_forward_hidden, dropout, device)