watch

Watch

Bases: TensorOp

Watch one or more tensors for later gradient computation.

Parameters:

Name	Type	Description	Default
inputs	Union[None, str, Iterable[str]]	which tensors to watch during future computation.	required
mode	Union[None, str, Iterable[str]]	What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument like "!infer" or "!train".	None

Source code in fastestimator\fastestimator\op\tensorop\gradient\watch.py

class Watch(TensorOp):
    """Watch one or more tensors for later gradient computation.

    Args:
        inputs: which tensors to watch during future computation.
        mode: What mode(s) to execute this Op in. For example, "train", "eval", "test", or "infer". To execute
            regardless of mode, pass None. To execute in all modes except for a particular one, you can pass an argument
            like "!infer" or "!train".
    """
    def __init__(self, inputs: Union[None, str, Iterable[str]], mode: Union[None, str, Iterable[str]] = None) -> None:
        super().__init__(inputs=inputs, outputs=inputs, mode=mode)
        self.in_list, self.out_list = True, True

    def forward(self, data: List[Tensor], state: Dict[str, Any]) -> List[Tensor]:
        for idx, tensor in enumerate(data):
            data[idx] = watch(tensor=tensor, tape=state['tape'])
        return data