code: def _encoder_forward(
      encoder: ViTEncoder,
      start_idx: int,
      end_idx: int,
      hidden_states: torch.Tensor,
      head_mask: Optional[torch.Tensor] = None,
      return_dict: bool = True,
      stage_manager: PipelineStageManager = None,
  ) -> Union[tuple, BaseModelOutput]:
      for i in range(start_idx, end_idx):
          layer_module = encoder.layer[i]
  
          layer_head_mask = (
              head_mask[i] if head_mask is not None else None
          )
  
          if encoder.gradient_checkpointing and encoder.training:
  
              def create_custom_forward(module):
                  def custom_forward(*inputs):
                      return module(*inputs, False)
  
                  return custom_forward
  
              layer_outputs = torch.utils.checkpoint.checkpoint(
                  create_custom_forward(layer_module),
                  hidden_states,
                  layer_head_mask,
              )
          else:
              layer_outputs = layer_module(
                  hidden_states, layer_head_mask, False
              )
  
          hidden_states = layer_outputs[0]
      if not stage_manager.is_last_stage():
          return hidden_states
      else:
          if not return_dict:
              return tuple(hidden_states)
          return BaseModelOutput(
              last_hidden_state=hidden_states,
              hidden_states=None,
              attentions=None,
          )
  
file_path: ColossalAI/colossalai/shardformer/modeling/vit.py
imports: import math
  from typing import List, Optional, Tuple, Union
  import torch
  from transformers.models.vit.modeling_vit import BaseModelOutput, ViTEncoder
  from transformers.utils import logging
  from colossalai.pipeline.stage_manager import PipelineStageManager
label_desc: Use image node_id 1 for calling a global function with example usage: _encoder_forward(encoder, start_idx, end_idx, hidden_states, head_mask, return_dict, stage_manager) and returns: hidden_states, BaseModelOutput, tuple, custom_forward, module