Method __init__

axlearn/common/trainer_config_modifier.py:137–140

(self, cfg: Config)

        mesh_shape: Required[Union[MeshShape, HybridMeshShape]] = REQUIRED

    def __init__(self, cfg: Config):
        super().__init__(cfg)
        cfg = self.config
        self._mesh_shape = cfg.mesh_shape

    def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
        """Overwrite the mesh shape.