hot fix bug

TomFrederik · Mar 16, 2022 · 0a12287 · 0a12287
1 parent 20a8070
commit 0a12287
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/unseal/__init__.py b/unseal/__init__.py
@@ -1,3 +1,3 @@
 from . import *
 
-__version__ = '0.1.6'
+__version__ = '0.1.7'
diff --git a/unseal/visuals/utils.py b/unseal/visuals/utils.py
@@ -117,6 +117,7 @@ def wrap_gpt_attn(
     attn_suffix: Optional[str] = None,
     layer_key_prefix: Optional[str] = None,
     out_proj_name: Optional[str] = 'out_proj',
+    batch_size: Optional[int] = None,
 ) -> Callable:
     # parse inputs
     if layer_key_prefix is None:
@@ -132,7 +133,8 @@ def wrap_gpt_attn(
         model.save_ctx[f'logit_layer_{layer}'], 
         model.layers[out_proj_name].weight,
         model.layers[unembedding_key].weight.T,
-        target_ids
+        target_ids,
+        batch_size,
     )
 
     return old_fn