docstr: Args:
      src: Tensor, with size [seq_len, batch_size]
      src_mask: Tensor, with size [seq_len, seq_len]
  
  Returns:
      output: Tensor, with size [seq_len, batch_size, n_token]
docstr_tok: ["Args",":","src",":","Tensor",",","with","size","[","seq_len",",","batch_size","]","src_mask",":","Tensor",",","with","size","[","seq_len",",","seq_len","]","Returns",":","output",":","Tensor",",","with","size","[","seq_len",",","batch_size",",","n_token","]"]
code: def forward(self, src, src_mask):
      """
      Args:
          src: Tensor, with size [seq_len, batch_size]
          src_mask: Tensor, with size [seq_len, seq_len]
  
      Returns:
          output: Tensor, with size [seq_len, batch_size, n_token]
      """
      src = self.embeddings(src) * math.sqrt(self.d_model)
      src = self.position(src)
      output = self.encoder(src, src_mask)
      output = self.decoder(output)
      return output
  
imports: import torch
  import torch.nn.functional
  import nni.retiarii.nn.pytorch
  from nni.retiarii import model_wrapper
  import nni
  import nni.retiarii.strategy
  from nni.retiarii.evaluator import FunctionalEvaluator
  from nni.retiarii.experiment.pytorch import RetiariiExperiment, RetiariiExeConfig
  import math
  from torchtext.datasets import WikiText2
  from torchtext.data.utils import get_tokenizer
  from torchtext.vocab import build_vocab_from_iterator
label_desc: Use image node_id 2 for calling the Transformer obj's underlying member method code with example usage: obj.forward(src, src_mask) and returns: output