Method _get_optimizer_model_params

axlearn/common/learner.py:208–211

(self, model_params: Nested[OptParam])

            )
        # Using self.optimizer: UpdateTransformation doesn't seem to work here.
        self.optimizer = cast(UpdateTransformation, self._add_child("optimizer", optimizer))
        if cfg.ema.decay is not None:
            self.ema: PartitionedGradientTransformation = cfg.ema.instantiate()
        if cfg.forward_fn_transformation is not None:
            self._forward_fn_transformation = maybe_instantiate(cfg.forward_fn_transformation)
        else:
            self._forward_fn_transformation = lambda fn: fn