From 6fed0032f238ba05b406e225afc204abed2a75b4 Mon Sep 17 00:00:00 2001
From: wassname <1103714+wassname@users.noreply.github.com>
Date: Wed, 12 Mar 2025 15:30:33 +0800
Subject: [PATCH] change default

---
 activation_store/collect.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/activation_store/collect.py b/activation_store/collect.py
index 52690f8..4393fa9 100644
--- a/activation_store/collect.py
+++ b/activation_store/collect.py
@@ -17,7 +17,7 @@ from torch import Tensor
 default_output_folder = (Path(__file__).parent.parent / "outputs").resolve()
 
 @torch.no_grad()
-def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutput, model: AutoModelForCausalLM, last_token=False, dtype=torch.float16) -> Dict[str, Tensor]:
+def default_postprocess_result(input: dict, trace: TraceDict, output: ModelOutput, model: AutoModelForCausalLM, last_token=True, dtype=torch.float16) -> Dict[str, Tensor]:
     """Make your own. This adds activations to output, and rearranges hidden states.
 
     Note the parquet write support float16, so we use that. It does not support float8, bfloat16, etc.