Merge pull request #116 from pamparamm/pooled_none_check

Partially resolve #115
2025-06-07 13:30:01 +03:00
parent d262a7dc7a 6c319ad5b4
commit 2437cd4daf
1 changed files with 1 additions and 1 deletions
@@ -117,7 +117,7 @@ def down_weight(tokens, weights, word_ids, base_emb, pooled_base, max_length, en
    w_mix = torch.tensor(w_mix, dtype=embs.dtype, device=embs.device).reshape((-1, 1, 1))

    weighted_emb = (w_mix * embs).sum(axis=0, keepdim=True)
-    if pooled and max_length:
+    if pooled is not None and max_length:
        pooled = weighted_emb[0, max_length - 1 : max_length, :]
    return weighted_emb, masked_current, pooled