code: def forward(
      self,
      tokens: torch.Tensor,
      segment_labels: torch.Tensor = None,
      positions: torch.Tensor = None,
      token_embeddings: torch.Tensor = None,
      attn_mask: torch.Tensor = None,
  ):
      assert self.has_segment_labels == (segment_labels is not None)
      assert self.has_positions == (positions is not None)
  
      traced_encoder_inputs = self._prepare_inputs(
          tokens, segment_labels, positions, token_embeddings, attn_mask
      )
      self.iter_ += 1
      if self.iter_ % 100 == 0:
          print("Iter: ", self.iter_)
          with torch.autograd.profiler.profile(
              enabled=True, use_cuda=True, record_shapes=True
          ) as prof:
              encoded_layers, pooled_output = self.traced_encoder(
                  *traced_encoder_inputs
              )
          print(
              prof.key_averages(group_by_input_shape=True).table(
                  sort_by="cuda_time"
              )
          )
      else:
          encoded_layers, pooled_output = self.traced_encoder(
              *traced_encoder_inputs
          )
  
      encoded_layers = list(torch.unbind(encoded_layers))
      return encoded_layers, pooled_output
  
file_path: pytext/pytext/models/representations/traced_transformer_encoder.py
imports: from typing import Tuple
  import torch
  import torch.cuda
  import torch.nn
  from fairseq.modules import TransformerSentenceEncoder
  from pytext.utils.usage import log_class_usage
label_desc: Use image node_id 2 for calling the TracedTransformerEncoder obj's underlying member method code with example usage: obj.forward(tokens, segment_labels, positions, token_embeddings, attn_mask) and returns: encoded_layers,  pooled_output