osc.models.models.forward_with_attns

forward_with_attns(model)

Context manager to set and remove hooks that collect attention maps.

Example

Usage:

>>> with forward_with_attns(model) as attns:
>>>     outputs = model(inputs)
>>> print(attns.items())

Parameters

model (Model) – Model where the attention hooks will be set.

Yields

A dictionary that will be populated after the forward pass.

Return type

Generator[Dict[str, Tensor], None, None]