class_docstr: Override TransformerEncoder to return the self-attn weights.
class_docstr_tok: ["Override","TransformerEncoder","to","return","the","self-attn","weights","."]
docstr: Override to return attention weights.
docstr_tok: ["Override","to","return","attention","weights","."]
code: def _apply_model_parallel(
      self, tensor, mask, **kwargs
  ) -> Tuple[torch.Tensor, torch.Tensor]:
      """
      Override to return attention weights.
      """
      chunks = PipelineHelper.split((tensor, mask))
      work_items = PipelineHelper.schedule_work_items(
          self.layers, chunks
      )
  
      for chunk_idx, layer_nos, next_device in work_items:
          s_weights = None
          try:
              s_tensor, s_mask = chunks[chunk_idx]
          except ValueError:
              s_tensor, s_mask, s_weights = chunks[chunk_idx]
          for layer_no in layer_nos:
              s_tensor, s_weights = self.layers[layer_no](
                  s_tensor, s_mask, **kwargs
              )
          chunks[chunk_idx] = PipelineHelper.chunk_to(
              (s_tensor, s_mask, s_weights), next_device
          )
      joined = PipelineHelper.join(chunks)
      tensor_out, out_mask, weights = joined
      return tensor_out, 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 3 for calling the TransformerReturnWeightsEncoder obj's underlying member method code with example usage: obj._apply_model_parallel(tensor, mask) and returns: tensor_out,  weights