fix trainable and requires_grad kwargs (#2932)

allenai · Jun 7, 2019 · 9a0e01f · 9a0e01f
1 parent 5f37783
commit 9a0e01f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/allennlp/models/bert_for_classification.py b/allennlp/models/bert_for_classification.py
@@ -59,7 +59,8 @@ def __init__(self,
         else:
             self.bert_model = bert_model
 
-        self.bert_model.requires_grad = trainable
+        for param in self.bert_model.parameters():
+            param.requires_grad = trainable
 
         in_features = self.bert_model.config.hidden_size
 

diff --git a/allennlp/modules/seq2vec_encoders/bert_pooler.py b/allennlp/modules/seq2vec_encoders/bert_pooler.py
@@ -39,7 +39,8 @@ def __init__(self, pretrained_model: Union[str, BertModel], requires_grad: bool
             model = pretrained_model
 
         self.pooler = model.pooler
-        self.pooler.requires_grad = requires_grad
+        for param in self.pooler.parameters():
+            param.requires_grad = requires_grad
         self._embedding_dim = model.config.hidden_size
 
     @overrides