didn't exist in pt-1.9 (#17644)

df1ec6b1 · Stas Bekman · GitHub · fba0b6a8 · df1ec6b1
Unverified Commit df1ec6b1 authored Jun 09, 2022 by Stas Bekman Committed by GitHub Jun 09, 2022
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

src/transformers/utils/import_utils.py src/transformers/utils/import_utils.py +4 -1

No files found.
--- a/src/transformers/utils/import_utils.py
+++ b/src/transformers/utils/import_utils.py
@@ -305,7 +305,10 @@ def is_torch_bf16_available():
        is_torch_gpu_bf16_available = False

    # checking CPU
-    if not hasattr(torch.cpu.amp, "autocast"):
+    try:
+        # multiple levels of AttributeError depending on the pytorch version so do them all in one check
+        _ = torch.cpu.amp.autocast
+    except AttributeError:
        is_torch_cpu_bf16_available = False

    return is_torch_cpu_bf16_available or is_torch_gpu_bf16_available