From 2f34bcf3e7b7f53181367faf7f492905612d205e Mon Sep 17 00:00:00 2001 From: fteufel <56223326+fteufel@users.noreply.github.com> Date: Mon, 12 Oct 2020 10:10:17 +0200 Subject: [PATCH] check for tpu availability in save_pretrained (#7699) Added is_torch_tpu_available() to the condition for saving a model as xla model. "xla_device" property of config can also be True on a non-xla device, when loading a checkpointthat was trained on xla before. Resolves #7695 --- src/transformers/modeling_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py index c152d76822..009fdc86cb 100755 --- a/src/transformers/modeling_utils.py +++ b/src/transformers/modeling_utils.py @@ -716,7 +716,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin): # If we save using the predefined names, we can load using `from_pretrained` output_model_file = os.path.join(save_directory, WEIGHTS_NAME) - if getattr(self.config, "xla_device", False): + if getattr(self.config, "xla_device", False) and is_torch_tpu_available(): import torch_xla.core.xla_model as xm if xm.is_master_ordinal():