class_docstr: The overall arch described in the BERT4Rec paper: (https://arxiv.org/abs/1904.06690)
  the encoder_layer was described in the section of 3.3, the output_layer was described in the
  section of 3.5
  
  Args:
      vocab_size (int): the item count including mask and padding
      max_len (int): the max length
      emb_dim (int): embedding dimension
      nhead (int): number of the transformation headers
      num_layers (int): number of the transformation layers
      dropout (float): the dropout probability
      device: (Optional[torch.device]).
  
  Example::
  
      input_kjt = KeyedJaggedTensor.from_lengths_sync(
          keys=["item"],
          values=torch.tensor([2, 4, 3, 4, 5]),
          lengths=torch.tensor([2, 3]),
      )
      device = (
          torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
      )
      input_kjt = input_kjt.to(device)
      bert4rec = BERT4Rec(
          vocab_size=6, max_len=3, emb_dim=4, nhead=4, num_layers=4, device=device
      )
      logits = bert4rec(input_kjt)
class_docstr_tok: ["The","overall","arch","described","in","the","BERT4Rec","paper",":","(","https",":","\/\/arxiv.org\/abs\/1904.06690",")","the","encoder_layer","was","described","in","the","section","of","3.3",",","the","output_layer","was","described","in","the","section","of","3.5","Args",":","vocab_size","(","int",")",":","the","item","count","including","mask","and","padding","max_len","(","int",")",":","the","max","length","emb_dim","(","int",")",":","embedding","dimension","nhead","(","int",")",":","number","of","the","transformation","headers","num_layers","(","int",")",":","number","of","the","transformation","layers","dropout","(","float",")",":","the","dropout","probability","device",":","(","Optional","[","torch.device","]",")",".","Example",":",":","input_kjt","=","KeyedJaggedTensor.from_lengths_sync","(","keys=","[","``","item","''","]",",","values=torch.tensor","(","[","2",",","4",",","3",",","4",",","5","]",")",",","lengths=torch.tensor","(","[","2",",","3","]",")",",",")","device","=","(","torch.device","(","``","cuda","''",")","if","torch.cuda.is_available","(",")","else","torch.device","(","``","cpu","''",")",")","input_kjt","=","input_kjt.to","(","device",")","bert4rec","=","BERT4Rec","(","vocab_size=6",",","max_len=3",",","emb_dim=4",",","nhead=4",",","num_layers=4",",","device=device",")","logits","=","bert4rec","(","input_kjt",")"]
docstr: forward function: first get the item embedding result and
  fit into transformer blocks and fit into the last linaer
  layer to get the final output
  
  Args:
      input (KeyedJaggedTensor): the input KeyedJaggedTensor
  
  Returns:
      torch.Tensor.
docstr_tok: ["forward","function",":","first","get","the","item","embedding","result","and","fit","into","transformer","blocks","and","fit","into","the","last","linaer","layer","to","get","the","final","output","Args",":","input","(","KeyedJaggedTensor",")",":","the","input","KeyedJaggedTensor","Returns",":","torch.Tensor","."]
code: def forward(self, input: KeyedJaggedTensor) -> torch.Tensor:
      """
      forward function: first get the item embedding result and
      fit into transformer blocks and fit into the last linaer
      layer to get the final output
  
      Args:
          input (KeyedJaggedTensor): the input KeyedJaggedTensor
  
      Returns:
          torch.Tensor.
      """
      dense_tensor = input["item"].to_padded_dense(
          desired_length=self.max_len
      )
      mask = (
          (dense_tensor > 0)
          .unsqueeze(1)
          .repeat(1, dense_tensor.size(1), 1)
          .unsqueeze(1)
      )
      # embedding the indexed sequence to sequence of vectors
      x = self.history(input)
      # running over multiple transformer blocks
      for transformer in self.transformer_blocks:
          x = transformer.forward(x, mask)
  
      return self.out(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 BERT4Rec obj's underlying member method code with example usage: obj.forward(input) and returns: self