missing crf output update (thanks Luca)

kermitt2 · Jul 12, 2023 · bc19e28 · bc19e28
1 parent 30e87a6
commit bc19e28
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/delft/sequenceLabelling/models.py b/delft/sequenceLabelling/models.py
@@ -1026,7 +1026,7 @@ def __init__(self, config, ntags=None, load_pretrained_weights=True, local_path:
         self.crf = ChainCRF()
         pred = self.crf(x)
 
-        self.model  = Model(inputs=[input_ids_in, features_input, token_type_ids, attention_mask], outputs=[x])
+        self.model  = Model(inputs=[input_ids_in, features_input, token_type_ids, attention_mask], outputs=[pred])
         self.config = config
 
     def get_generator(self):