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",")","]",")"]
docstr: forward function
  
  Args:
      x (torch.Tensor): the input tensor
      mask (torch.BoolTensor): determine which position has been masked
  
  Returns:
      torch.Tensor.
docstr_tok: ["forward","function","Args",":","x","(","torch.Tensor",")",":","the","input","tensor","mask","(","torch.BoolTensor",")",":","determine","which","position","has","been","masked","Returns",":","torch.Tensor","."]
code: def forward(
      self, x: torch.Tensor, mask: torch.BoolTensor
  ) -> torch.Tensor:
      """
      forward function
  
      Args:
          x (torch.Tensor): the input tensor
          mask (torch.BoolTensor): determine which position has been masked
  
      Returns:
          torch.Tensor.
      """
      x = self.input_sublayer(
          x, lambda _x: self.attention.forward(_x, _x, _x, mask=mask)
      )
      x = self.output_sublayer(x, self.feed_forward)
      return self.dropout(x)
  
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 2 for calling the TransformerBlock obj's underlying member method code with example usage: obj.forward(x, mask) and returns: self