riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-07-03 12:50:06 +06:00
Code Issues Actions 28 Packages Projects Releases Wiki Activity

fix: bf16 with TPU is allowed in configuration (#38670)

Browse Source 
* fix: tpu bf16

* fix: style

---------

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
This commit is contained in:
ye 2025-06-11 21:35:01 +09:00 committed by GitHub
parent c8c1e525ed
commit b426c2b313
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/training_args.py
View File

@ -1692,7 +1692,7 @@ class TrainingArguments:
# cpu
raise ValueError("Your setup doesn't support bf16/(cpu, tpu, neuroncore). You need torch>=1.10")
elif not self.use_cpu:
if not is_torch_bf16_gpu_available():
if not is_torch_bf16_gpu_available() and not is_torch_xla_available(): # added for tpu support
error_message = "Your setup doesn't support bf16/gpu."
if is_torch_cuda_available():
error_message += " You need Ampere+ GPU with cuda>=11.0"