Function make_hook

GQLA_preprint/src/compression.py:126–132

(li)

    final_hiddens: list[torch.Tensor] | None = [] if capture_final else None

    def make_hook(li):
        def _hook(_m, args, kwargs):
            x = kwargs.get("hidden_states") if kwargs else None
            if x is None and args:
                x = args[0]
            captured[li].append(x.detach().to("cpu"))
        return _hook

    def _final_hook(_m, args, kwargs):
        x = kwargs.get("hidden_states") if kwargs else None