SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Should check that torch TPU is available (#5636)

Browse Source
This commit is contained in:
Lysandre Debut
2020-07-09 13:54:32 -04:00
committed by GitHub
parent 3cc23eee06
commit b25f7802de
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/transformers/modeling_utils.py
View File

@@ -34,6 +34,7 @@ from .file_utils import (
cached_path, cached_path,
hf_bucket_url, hf_bucket_url,
is_remote_url, is_remote_url,
is_torch_tpu_available,
) )
from .generation_utils import GenerationMixin from .generation_utils import GenerationMixin
@@ -794,7 +795,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
} }
return model, loading_info return model, loading_info
if hasattr(config, "xla_device") and config.xla_device: if hasattr(config, "xla_device") and config.xla_device and is_torch_tpu_available():
import torch_xla.core.xla_model as xm import torch_xla.core.xla_model as xm
model = xm.send_cpu_data_to_device(model, xm.xla_device()) model = xm.send_cpu_data_to_device(model, xm.xla_device())