Accept tensorflow-rocm package when checking TF availability (#11595)

864c1dfe · Mats Sjöberg · GitHub · 3e3e41ae · 864c1dfe
Unverified Commit 864c1dfe authored May 05, 2021 by Mats Sjöberg Committed by GitHub May 05, 2021
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/file_utils.py src/transformers/file_utils.py +1 -0

No files found.
--- a/src/transformers/file_utils.py
+++ b/src/transformers/file_utils.py
@@ -88,6 +88,7 @@ if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VA
            "tf-nightly-cpu",
            "tf-nightly-gpu",
            "intel-tensorflow",
+            "tensorflow-rocm",
        )
        _tf_version = None
        # For the metadata, we have to look for both tensorflow and tensorflow-cpu