fix distilbert

Tiiiger · Dec 22, 2019 · e18bd86 · e18bd86
1 parent fc56310
commit e18bd86
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bert_score/utils.py b/bert_score/utils.py
@@ -108,7 +108,7 @@ def get_model(model_type, num_layers, all_layers=None):
                     torch.nn.ModuleList([layer for layer in model.encoder.layer[:num_layers]])
         elif hasattr(model, 'transformer'): # bert, roberta
             model.transformer.layer =\
-                torch.nn.ModuleList([layer for layer in model.encoder.layer[:num_layers]])
+                torch.nn.ModuleList([layer for layer in model.transformer.layer[:num_layers]])
         else:
             raise ValueError("Not supported")
     else: