Skip to content

Commit 350a1b1

Browse files
Sid Mohanclaude
andcommitted
Fix CRF label_weights: use register_buffer in __init__
Plain attribute assignment in __init__ conflicts with register_buffer() in set_label_weights() — PyTorch raises KeyError because the attribute already exists but isn't a buffer. Initialize as a None buffer instead. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent e2acf18 commit 350a1b1

1 file changed

Lines changed: 2 additions & 1 deletion

File tree

  • pii-ner-v1/src/datafog_pii_ner/model

pii-ner-v1/src/datafog_pii_ner/model/crf.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@ def __init__(self, hidden_dim: int, num_labels: int, dropout: float = 0.1):
2121
self.num_labels = num_labels
2222
# label_weights: (num_labels,) tensor mapping label_id -> sequence weight.
2323
# Set via set_label_weights(). None means uniform weighting.
24-
self.label_weights: torch.Tensor | None = None
24+
# Use register_buffer so set_label_weights can update it later.
25+
self.register_buffer("label_weights", None)
2526

2627
def set_label_weights(self, weights: torch.Tensor):
2728
"""Set per-label weights for tier-weighted sequence loss.

0 commit comments

Comments
 (0)