code: def main():
      colossalai.launch_from_torch(config="./config.py")
  
      logger = get_dist_logger()
  
      # trace the model with meta data
      model = resnet50(num_classes=10).cuda()
  
      input_sample = {
          "x": torch.rand(
              [
                  gpc.config.BATCH_SIZE
                  * torch.distributed.get_world_size(),
                  3,
                  32,
                  32,
              ]
          ).to("meta")
      }
      device_mesh = DeviceMesh(
          physical_mesh_id=torch.tensor([0, 1, 2, 3]),
          mesh_shape=[2, 2],
          init_process_group=True,
      )
      model, solution = initialize_model(
          model,
          input_sample,
          device_mesh=device_mesh,
          return_solution=True,
      )
  
      if gpc.get_global_rank() == 0:
          for node_strategy in solution:
              print(node_strategy)
      # build criterion
      criterion = torch.nn.CrossEntropyLoss()
  
      # optimizer
      optimizer = torch.optim.SGD(
          model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4
      )
  
      # lr_scheduler
      lr_scheduler = CosineAnnealingLR(
          optimizer, total_steps=gpc.config.NUM_EPOCHS
      )
  
      for epoch in range(gpc.config.NUM_EPOCHS):
          model.train()
  
          # if we use synthetic data
          # we assume it only has 10 steps per epoch
          num_steps = range(10)
          progress = tqdm(num_steps)
  
          for _ in progress:
              # generate fake data
              img, label = synthesize_data()
  
              img = img.cuda()
              label = label.cuda()
              optimizer.zero_grad()
              output = model(img)
              train_loss = criterion(output, label)
              train_loss.backward(train_loss)
              torch.cuda.synchronize()
              optimizer.step()
          lr_scheduler.step()
  
          # run evaluation
          model.eval()
          correct = 0
          total = 0
  
          # if we use synthetic data
          # we assume it only has 10 steps for evaluation
          num_steps = range(10)
          progress = tqdm(num_steps)
  
          for _ in progress:
              # generate fake data
              img, label = synthesize_data()
  
              img = img.cuda()
              label = label.cuda()
  
              with torch.no_grad():
                  output = model(img)
                  test_loss = criterion(output, label)
              pred = torch.argmax(output, dim=-1)
              correct += torch.sum(pred == label)
              total += img.size(0)
  
          logger.info(
              f"Epoch {epoch} - train loss: {train_loss:.5}, test loss: {test_loss:.5}, acc: {correct / total:.5}, lr: {lr_scheduler.get_last_lr()[0]:.5g}",
              ranks=[0],
          )
  
file_path: ColossalAI/examples/tutorial/auto_parallel/auto_parallel_with_resnet.py
imports: import torch
  from torchvision.models import resnet50
  from tqdm import tqdm
  import colossalai
  from colossalai.auto_parallel.tensor_shard.initialize import initialize_model
  from colossalai.device.device_mesh import DeviceMesh
  from colossalai.legacy.core import global_context
  from colossalai.logging import get_dist_logger
  from colossalai.nn.lr_scheduler import CosineAnnealingLR
label_desc: Use image node_id 2 for calling a global function with example usage: main() without return types