Fix TensorFlow pakage check (huggingface#26842)

Add tf-nightly-rocm to _is_tf_available check
EduardoPach · Nov 19, 2023 · a3fb1de · a3fb1de
1 parent 7ca14a1
commit a3fb1de
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py
@@ -168,6 +168,7 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> Union[
                 "tf-nightly",
                 "tf-nightly-cpu",
                 "tf-nightly-gpu",
+                "tf-nightly-rocm",
                 "intel-tensorflow",
                 "intel-tensorflow-avx512",
                 "tensorflow-rocm",