From 756b6539138d00446d973e63640727bcc868f6b6 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Wed, 12 Mar 2025 14:06:40 +0800 Subject: [PATCH] handle 0d tensors like loss --- activation_store/collect.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/activation_store/collect.py b/activation_store/collect.py index 2256252..52690f8 100644 --- a/activation_store/collect.py +++ b/activation_store/collect.py @@ -42,8 +42,12 @@ def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutpu o['attention_mask'] = input['attention_mask'] if 'label' in input: o['label'] = input['label'] - input = output = acts = None + # convert any 0d tensors like loss to 1d, by repeating along batch dimension + for k, v in o.items(): + if v.dim() == 0: + bs = input['input_ids'].shape[0] + o[k] = v.repeat(bs) return o @@ -105,7 +109,7 @@ def activation_store(loader: DataLoader, model: AutoModelForCausalLM, dataset_na for bo in iterator: bs = len(next(iter(bo.values()))) - assert all(len(v) == bs for v in bo.values()), "must return Dict[str,Tensor] and all tensors with same batch size a first dimension" + assert all(len(v) == bs for v in bo.values()), "must return Dict[str,Tensor] and all tensors with same batch size as first dimension" # or maybe better compression to `writer.write(example, key)` for each writer.write_batch(bo)