fix roberta conversion script

2019-09-24 09:48:23 +02:00
parent 2b11fa5174
commit 8ba44ced95
2 changed files with 3 additions and 3 deletions
--- a/pytorch_transformers/init.py
+++ b/pytorch_transformers/init.py
@@ -138,7 +138,7 @@ if _tf_available:
                                  TF_XLM_PRETRAINED_MODEL_ARCHIVE_MAP)

    from .modeling_tf_roberta import (TFRobertaPreTrainedModel, TFRobertaMainLayer,
-                                      TFRobertaModel, TFRobertaLMHead,
+                                      TFRobertaModel, TFRobertaForMaskedLM,
                                      TFRobertaForSequenceClassification,
                                      load_roberta_pt_weights_in_tf2,
                                      TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP)