Method __init__

axlearn/common/trainer_config_modifier.py:279–281

(self, cfg: Config)

        partition_specs: Required[Dict[str, PartitionSpec]] = REQUIRED

    def __init__(self, cfg: Config):
        super().__init__(cfg)
        self._attribute_dicts = self.config.partition_specs

    def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
        """Update the partition_spec attributes for the specified modules.