code: def _register_colo_modules(self):
      from colossalai.legacy.nn.parallel.layers import (
          ColoEmbedding,
          ColoLinear,
          register_colo_module,
      )
  
      register_colo_module(torch.nn.Linear, ColoLinear())
      register_colo_module(torch.nn.Embedding, ColoEmbedding())
  
file_path: colossalai/colossalai/legacy/zero/gemini/colo_init_context.py
imports: from typing import Any, Iterator, Optional, Tuple, Union
  import torch
  from torch import nn
  from colossalai.legacy.tensor import ProcessGroup
  from colossalai.tensor import ColoParameter, ColoTensor
  from colossalai.utils.model.utils import InsertPostInitMethodToModuleSubClasses
label_desc: Use image node_id 2 for calling the ColoInitContext obj's underlying member method code with example usage: obj._register_colo_modules() without return types