Update transformer.py

som-shahlab · Dec 10, 2023 · 746cb8e · 746cb8e
1 parent 431ef88
commit 746cb8e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/femr/models/transformer.py b/src/femr/models/transformer.py
@@ -169,7 +169,7 @@ def forward(self, features: torch.Tensor, batch: Mapping[str, torch.Tensor]):
 
 
 class CLMBRTaskHead(nn.Module):
-    def __init__(self, hidden_size: int, clmbr_vocab_size: int):
+    def __init__(self, hidden_size: int, clmbr_vocab_size: int, **kwargs):
         super().__init__()
 
         self.final_layer = nn.Linear(hidden_size, clmbr_vocab_size)