Add attention_mask argument to loss_fn() and lm_cross_entropy_loss() …

…and adjust the cross entropy calculation to ignore masked (padding) tokens.
TransformerLensOrg · Oct 18, 2024 · 7eb5679 · 7eb5679
1 parent 15ae297
commit 7eb5679
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -3,7 +3,7 @@
     description="An implementation of transformers tailored for mechanistic interpretability."
     license="MIT"
     name="transformer-lens"
-    packages=[{include="transformer_lens"}]
+    packages=[{include="transformer_lens"}, {include="transformer_lens/py.typed"}]
     readme="README.md"
     # Version is automatically set by the pipeline on release
     version="0.0.0"

diff --git a/transformer_lens/py.typed b/transformer_lens/py.typed