code: def forward(
      self, hidden_states: Optional[Tuple[torch.FloatTensor]]
  ) -> torch.FloatTensor:
      hidden_states = self.c_fc(hidden_states)
      hidden_states = checkpoint(self.c_proj, hidden_states)
      hidden_states = self.act(hidden_states)
  
      return hidden_states
  
file_path: ColossalAI/tests/test_auto_parallel/test_tensor_shard/test_checkpoint.py
imports: from typing import Optional, Tuple
  import pytest
  import torch
  import torch.nn
  from torch.utils.checkpoint import checkpoint
  from transformers.pytorch_utils import Conv1D
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.initialize import launch
  from colossalai.logging import disable_existing_loggers
  from colossalai.testing import rerun_if_address_is_in_use, run_on_environment_flag, spawn
label_desc: Use image node_id 2 for calling the GPT2MLPWithCkpt obj's underlying member method code with example usage: obj.forward(hidden_states) and returns: hidden_states