class_docstr: This corresponds to the EncoderBlock class in the scenic/vivit implementation.
class_docstr_tok: ["This","corresponds","to","the","EncoderBlock","class","in","the","scenic\/vivit","implementation","."]
code: def forward(
      self, hidden_states, head_mask=None, output_attentions=False
  ):
      self_attention_outputs = self.attention(
          # in Vivit, layernorm is applied before self-attention
          self.layernorm_before(hidden_states),
          head_mask,
          output_attentions=output_attentions,
      )
      attention_output = self_attention_outputs[0]
      # add self attentions if we output attention weights
      outputs = self_attention_outputs[1:]
  
      # first residual connection
      hidden_states = attention_output + hidden_states
  
      # in Vivit, layernorm is also applied after self-attention
      layer_output = self.layernorm_after(hidden_states)
      layer_output = self.intermediate(layer_output)
  
      # second residual connection is done here
      layer_output = self.output(layer_output, hidden_states)
  
      outputs = (layer_output,) + outputs
  
      return outputs
  
imports: import math
  from typing import Optional, Set, Tuple, Union
  import torch
  import torch.utils.checkpoint
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  from ...activations import ACT2FN
  from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ImageClassifierOutput
  from ...modeling_utils import PreTrainedModel
  from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
  from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from .configuration_vivit import VivitConfig
label_desc: Use image node_id 2 for calling the VivitLayer obj's underlying member method code with example usage: obj.forward(hidden_states, head_mask, output_attentions) and returns: outputs