riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-02 19:21:31 +06:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Fix gather for TPU (#13813)

Browse Source
This commit is contained in:
Sylvain Gugger 2021-09-30 11:32:40 -04:00 committed by GitHub
parent 7db2a79b38
commit 269c3d1400
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/trainer_pt_utils.py
View File

@ -152,6 +152,8 @@ def nested_xla_mesh_reduce(tensors, name):
if isinstance(tensors, (list, tuple)):
return type(tensors)(nested_xla_mesh_reduce(t, f"{name}_{i}") for i, t in enumerate(tensors))
if tensors.ndim == 0:
tensors = tensors[None]
return xm.mesh_reduce(name, tensors, torch.cat)
else:
raise ImportError("Torch xla must be installed to use `nested_xla_mesh_reduce`")