class_docstr: Override TransformerEncoder to return the self-attn weights.
class_docstr_tok: ["Override","TransformerEncoder","to","return","the","self-attn","weights","."]
docstr: Forward pass.
  
  Propagate kwargs
docstr_tok: ["Forward","pass",".","Propagate","kwargs"]
code: def forward(  # type: ignore
      self,
      input: torch.LongTensor,
      positions: Optional[torch.LongTensor] = None,
      segments: Optional[torch.LongTensor] = None,
      **kwargs,
  ) -> Union[
      Tuple[torch.Tensor, Optional[torch.Tensor]],
      Tuple[torch.Tensor, torch.BoolTensor, Optional[torch.Tensor]],
  ]:
      """
      Forward pass.
  
      Propagate kwargs
      """
      # embed input
      tensor, mask = self.forward_embedding(input, positions, segments)
  
      if self.variant == "xlm" or self.variant == "bart":
          tensor = self.norm_embeddings(tensor)
  
      # --dropout on the embeddings
      tensor = self.dropout(tensor)
  
      tensor *= mask.unsqueeze(-1).type_as(tensor)
  
      # apply transformer layers
      tensor = self.forward_layers(tensor, mask, **kwargs)
  
      ###################
      # BEGIN DIFFERENCE#
      ###################
      tensor, weights = tensor
      ###################
      # \End  DIFFERENCE#
      ###################
  
      if self.variant == "prelayernorm":
          tensor = self.norm_embeddings(tensor)
  
      # reduce output
      tensor, out_mask = self.reduce_output(tensor, mask)
      if out_mask is not None:
          return tensor, out_mask, weights
      else:
          return tensor, weights
  
imports: import torch
  import torch.nn
  import torch.nn.functional
  from typing import Optional, Tuple, Dict, Union
  from parlai.agents.transformer.modules import MultiHeadAttention, TransformerEncoder, TransformerEncoderLayer
  from parlai.agents.transformer.polyencoder import PolyencoderAgent, PolyEncoderModule, PolyBasicAttention
  from parlai.core.dict import DictionaryAgent
  from parlai.core.loader import register_agent
  from parlai.core.opt import Opt
  from parlai.core.params import ParlaiParser
  from parlai.core.torch_agent import Batch
  from parlai.utils.torch import PipelineHelper
label_desc: Use image node_id 1 for calling the TransformerReturnWeightsEncoder obj's underlying member method code with example usage: obj.forward(input, positions, segments) and returns: tensor,  out_mask,  weights, tensor,  weights