Method add_state_update

axlearn/common/module.py:1123–1125

Adds a state update to the output collection of the current context.

(self, name: str, value: Tensor)

        return self.get_invocation_context().state

    def add_summary(self, name: str, value: Union[Summable, Tensor, Summary]):
        """Adds the named value to `OutputCollection.summaries`.

        Args:
            name: The name of the item to add.
            value: The value to add.