From e099d03638322e73d874e336365bbb0c8054f98d Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Fri, 14 Mar 2025 17:09:37 +0800
Subject: [PATCH] wip

---
 activation_store/collect.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/activation_store/collect.py b/activation_store/collect.py
index d68d73a..77a3b30 100644
--- a/activation_store/collect.py
+++ b/activation_store/collect.py
@@ -36,11 +36,12 @@ def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutpu
             aas = [v.output[0] if isinstance(v.output, tuple) else v.output for k, v in trace.items() if k in group]
             assert len(aas) > 0, f"no activations found for {group}"
             aas = torch.stack([a[:, token_index].to(dtype) for a in aas], dim=1)
-            acts[k] = aas
+            acts[f'acts-{k}'] = aas
     else:
         acts = {f'act-{k}': 
-                v.output[0][:, token_index].to(dtype) if isinstance(v.output, tuple) else v.output[:, token_index]
+                v.output[0] if isinstance(v.output, tuple) else v.output
                 for k, v in trace.items()}
+        acts = {k: v[:, token_index].to(dtype) for k, v in acts.items() if v is not None}
     del trace
     
     # batch must be first, also the writer supports float16 so lets use that