update loading logics

2019-09-24 14:03:58 +02:00
parent 271f213621
commit f5397ffc3b
4 changed files with 55 additions and 35 deletions
--- a/pytorch_transformers/init.py
+++ b/pytorch_transformers/init.py
@@ -163,7 +163,7 @@ if _tf_available and _torch_available:
 # Files and general utilities
 from .file_utils import (PYTORCH_TRANSFORMERS_CACHE, PYTORCH_PRETRAINED_BERT_CACHE,
                         cached_path, add_start_docstrings, add_end_docstrings,
-                         WEIGHTS_NAME, TF_WEIGHTS_NAME, CONFIG_NAME)
+                         WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME, CONFIG_NAME)

 def is_torch_available():
    return _torch_available