Remove pooler from BERT for Masked LM architecture (guillaume-be#350)

SpirosMakris · Apr 10, 2023 · 2e5f4a9 · 2e5f4a9
1 parent 46fbf62
commit 2e5f4a9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -12,6 +12,7 @@ All notable changes to this project will be documented in this file. The format
 
 ## Fixed
 - MIN/MAX computation for float-like (was set to infinity instead of min/max)
+- Remove the (unused) pooler from the set of weights for BERT Masked LM architecture
 
 ## [0.20.0] - 2023-01-21
 ## Added

diff --git a/src/bert/bert_model.rs b/src/bert/bert_model.rs
@@ -559,7 +559,7 @@ impl BertForMaskedLM {
     {
         let p = p.borrow();
 
-        let bert = BertModel::new(p / "bert", config);
+        let bert = BertModel::new_with_optional_pooler(p / "bert", config, false);
         let cls = BertLMPredictionHead::new(p / "cls", config);
 
         BertForMaskedLM { bert, cls }