Method __call__

tests/hooks/test_group_offloading.py:137–140

(self, x: torch.Tensor)

        self.register_modules(model=model)

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        for _ in range(2):
            x = x + 0.1 * self.model(x)
        return x


class LayerOutputTrackerHook(ModelHook):