Using configuration for xla_device

[LysandreJik]

Using the model configuration file is usually preferred to adding new arguments on methods such as save_pretrained/from_pretrained.

There is no addition to the examples/run_glue_tpu.py file as the xla_device=True flag was already specified to the configuration file. That's all that's needed now 🙂

Here's a brief explanation of how that would work:

from transformers import BertConfig, BertModel

config = BertConfig.from_pretrained("bert-base-cased", xla_config=True)  # config.xla_device == True
config.save_pretrained("directory")
config = BertConfig.from_pretrained("directory")  # config.xla_device == True
config = BertConfig.from_pretrained("directory", xla_device=False)  # config.xla_device == False

Here's how you would save/load using an XLA device:

from transformers import BertConfig, BertModel

# With that configuration flag the model is loaded automatically on TPU
config = BertConfig.from_pretrained("bert-base-cased", xla_config=True)
model = BertModel.from_pretrained("bert-base-cased", config=config)

# Saved automatically on TPU
model.save_pretrained("directory")

# If we want to load without TPU
config = BertConfig.from_pretrained("directory", xla_config=False)
model = BertModel.from_pretrained("directory", config=config)

It can be done simply as well without explicitly using the config:

from transformers import BertModel

model = BertModel.from_pretrained("bert-base-cased", xla_device=True)
model.config.xla_device  # True

[jysohn23]

I like it better too. Thanks!