Fix the cuda-specific transformer utils for ROCm

57dea7f2 · hubertlu-tw · cb8b7a88 · 57dea7f2
Commit 57dea7f2 authored Aug 08, 2022 by hubertlu-tw
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

apex/transformer/testing/distributed_test_base.py apex/transformer/testing/distributed_test_base.py +4 -1

No files found.
--- a/apex/transformer/testing/distributed_test_base.py
+++ b/apex/transformer/testing/distributed_test_base.py
@@ -20,7 +20,10 @@ except ImportError:
 _TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION = Version("470.42.01")
 _driver_version = None
 if torch.cuda.is_available():
-    _driver_version = parse(collect_env.get_nvidia_driver_version(collect_env.run))
+    if collect_env.get_nvidia_driver_version(collect_env.run) != None:
+        _driver_version = parse(collect_env.get_nvidia_driver_version(collect_env.run))
+    else:
+        _driver_version = None
 HAS_TORCH_UCC_COMPAT_NVIDIA_DRIVER = _driver_version is not None and _driver_version >= _TORCH_UCC_COMPAT_NVIDIA_DRIVER_VERSION