class_docstr: Constructs attentions and saves their weights!
class_docstr_tok: ["Constructs","attentions","and","saves","their","weights","!"]
docstr: Override Polyencoder.encode to *only* return the coded/self-attn attention
  weights.
docstr_tok: ["Override","Polyencoder.encode","to","*","only","*","return","the","coded\/self-attn","attention","weights","."]
code: def encode(
      self,
      cand_tokens: Optional[torch.Tensor],
      **ctxt_inputs: torch.Tensor,
  ) -> Tuple[
      torch.Tensor,
      torch.Tensor,
      torch.BoolTensor,
      Optional[torch.Tensor],
  ]:
      """
      Override Polyencoder.encode to *only* return the coded/self-attn attention
      weights.
      """
      assert len(ctxt_inputs) > 0
      assert "ctxt_tokens" in ctxt_inputs
      assert len(ctxt_inputs["ctxt_tokens"].shape) == 2
      assert self.type == "codes"
  
      cand_embed = None
      if cand_tokens is not None:
          if len(cand_tokens.shape) != 3:
              cand_tokens = cand_tokens.unsqueeze(1)
          bsz = cand_tokens.size(0)
          num_cands = cand_tokens.size(1)
          cand_embed = self.encoder_cand(
              cand_tokens.view(bsz * num_cands, -1)
          )
          cand_embed = cand_embed.view(bsz, num_cands, -1)
  
      bsz = self._get_context_batch_size(**ctxt_inputs)
      # get context_representation. Now that depends on the cases.
      ctxt_out, ctxt_mask, ctxt_self_attn_weights = self.encoder_ctxt(
          **self._context_encoder_input(ctxt_inputs)
      )
      ctxt_self_attn_weights = F.softmax(ctxt_self_attn_weights, dim=-1)
      ctxt_rep, ctxt_code_weights = self.attend(
          self.code_attention,
          queries=self.codes.repeat(bsz, 1, 1),
          keys=ctxt_out,
          values=ctxt_out,
          mask=ctxt_mask,
      )
  
      return ctxt_rep, ctxt_code_weights, ctxt_mask, cand_embed
  
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 4 for calling the PolyencoderReturnWeightsModule obj's underlying member method code with example usage: obj.encode(cand_tokens) and returns: ctxt_rep,  ctxt_code_weights,  ctxt_mask,  cand_embed