Method __init__

axlearn/common/trainer_config_modifier.py:99–102

(self, cfg: Config)

        remat_policies: Required[Dict[str, RematSpec]] = REQUIRED

    def __init__(self, cfg: Config):
        super().__init__(cfg)
        cfg = self.config
        self._remat_policies = cfg.remat_policies

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