Is it possible to deal with the activations before exiting the tracing context?

[Shiroha-Offical]

I'm trying to intervene the modules in .generate call with some conditions. However, the conditions are calculated by the activations before exiting the tracing context. Therefore, I couldn't do anything without the value of the activations.

So, is there any solution or alternative to achieve it?
I will be grateful for any help you can provide.

Here are my desired behavior with demonstration in code:

with llm.generate(prompt_template(query), max_new_tokens=max_lenth, top_k=1) as tracer:
       all_tokens = llm.generator.output.detach().cpu().save()
       inputs = llm.lm_head.input.save()
       for count in range(intervene_steps):
            for layer in range(num_layers):
                current_head_wise_activations = llm.model.layers[layer].next().self_attn.o_proj.output.detach().cpu().save()
                # deal with the current_head_wise_activations