also gathering file names in file_utils

2019-09-05 02:34:09 +02:00
parent 121f88cae3
commit 85df4f7cca
4 changed files with 11 additions and 10 deletions
--- a/pytorch_transformers/file_utils.py
+++ b/pytorch_transformers/file_utils.py
@@ -48,6 +48,10 @@ except (AttributeError, ImportError):

 PYTORCH_TRANSFORMERS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility

+WEIGHTS_NAME = "pytorch_model.bin"
+TF_WEIGHTS_NAME = 'model.ckpt'
+CONFIG_NAME = "config.json"
+
 logger = logging.getLogger(__name__)  # pylint: disable=invalid-name

 if not six.PY2: