Fix tf v1 gpu import error (#267)

markurtz · web-flow · commit d72c73edfaec · 2021-06-02T11:58:50.000-04:00
diff --git a/integrations/pytorch/notebooks/sparse_quantized_transfer_learning.ipynb b/integrations/pytorch/notebooks/sparse_quantized_transfer_learning.ipynb
@@ -454,4 +454,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
diff --git a/src/sparseml/tensorflow_v1/base.py b/src/sparseml/tensorflow_v1/base.py
@@ -102,7 +102,13 @@ def check_tensorflow_install(
             raise tensorflow_err
         return False
 
-    return check_version("tensorflow", min_version, max_version, raise_on_error)
+    return check_version(
+        "tensorflow",
+        min_version,
+        max_version,
+        raise_on_error,
+        alternate_package_names=["tensorflow-gpu"],
+    )
 
 
 def check_tf2onnx_install(
