code: def __init__(self, dim: int = 256):
      super().__init__()
      intermediate_dim = dim * 4
      self.dense_1 = torch.nn.Linear(dim, intermediate_dim)
      self.activation = torch.nn.GELU()
      self.dense_2 = torch.nn.Linear(intermediate_dim, dim)
      self.dropout = torch.nn.Dropout(0.1)
  
file_path: colossalai/tests/test_legacy/test_pipeline/test_pipelinable.py
imports: import pytest
  import torch
  from colossalai.legacy.pipeline.pipelinable import PipelinableContext
  from colossalai.testing import rerun_if_address_is_in_use, spawn
label_desc: Use image node_id 1 to create a new MLP object from inherited base classes: torch.nn with example: obj = MLP(dim)