Method apply

composer/algorithms/sam/sam.py:167–180

(self, event: Event, state: State, logger: Optional[Logger])

        return event == Event.AFTER_LOAD

    def apply(self, event: Event, state: State, logger: Optional[Logger]) -> Optional[int]:
        assert state.optimizers is not None

        state.optimizers = tuple(
            SAMOptimizer(
                base_optimizer=optimizer,
                rho=self.rho,
                epsilon=self.epsilon,
                interval=self.interval,
            ) for optimizer in ensure_tuple(state.optimizers)
        )

        # Switch to ClosureGradScaler as SAM supports and requires it
        state.scaler = ClosureGradScaler()