need to freeze the model (NVIDIA#3540)

Signed-off-by: Yi Dong <[email protected]> Co-authored-by: Eric Harper <[email protected]>
bogdal1993 · Jan 28, 2022 · 8b2d27a · 8b2d27a
1 parent 101977e
commit 8b2d27a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py b/nemo/collections/nlp/models/text_classification/ptune_text_classification_model.py
@@ -87,8 +87,8 @@ def __init__(self, cfg: DictConfig, trainer: Trainer = None):
             trainer=trainer,
         )
 
-        for param in self.model.parameters():
-            param.requires_grad = cfg.use_lm_finetune
+        if not cfg.use_lm_finetune:
+            self.model.freeze()
 
         hidden_size = self.model.cfg.hidden_size