From 6c319ad5b480fe0b7f9d525b83501c7f3f7c8139 Mon Sep 17 00:00:00 2001 From: Pam Date: Sat, 7 Jun 2025 09:38:13 +0500 Subject: [PATCH] Partially resolve #115 --- prompt_control/adv_encode.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prompt_control/adv_encode.py b/prompt_control/adv_encode.py index de69d9e..7bc40b3 100644 --- a/prompt_control/adv_encode.py +++ b/prompt_control/adv_encode.py @@ -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