Interpretability with tensordict and torch hooks.