MCPcopy Create free account
hub / github.com/MuLabPKU/TransArch / _final_hook

Function _final_hook

GQLA_preprint/src/compression.py:134–138  ·  view source on GitHub ↗
(_m, args, kwargs)

Source from the content-addressed store, hash-verified

132 return _hook
133
134 def _final_hook(_m, args, kwargs):
135 x = kwargs.get("hidden_states") if kwargs else None
136 if x is None and args:
137 x = args[0]
138 final_hiddens.append(x.detach().to("cpu"))
139
140 hooks = [
141 layer.self_attn.register_forward_pre_hook(make_hook(i), with_kwargs=True)

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected