diff --git a/third_party/tsl/third_party/gpus/check_cuda_libs.py b/third_party/tsl/third_party/gpus/check_cuda_libs.py
index b1a10a86b9aac6..a1d47efcc93a81 100644
--- a/third_party/tsl/third_party/gpus/check_cuda_libs.py
+++ b/third_party/tsl/third_party/gpus/check_cuda_libs.py
@@ -27,16 +27,10 @@
 import os
 import os.path
 import platform
+import shutil
 import subprocess
 import sys
 
-# pylint: disable=g-import-not-at-top,g-importing-member
-try:
-  from shutil import which
-except ImportError:
-  from distutils.spawn import find_executable as which
-# pylint: enable=g-import-not-at-top,g-importing-member
-
 
 class ConfigError(Exception):
   pass
@@ -59,7 +53,7 @@ def check_cuda_lib(path, check_soname=True):
   """
   if not os.path.isfile(path):
     raise ConfigError("No library found under: " + path)
-  objdump = which("objdump")
+  objdump = shutil.which("objdump")
   if check_soname and objdump is not None and not _is_windows():
     # Decode is necessary as in py3 the return type changed from str to bytes
     output = subprocess.check_output([objdump, "-p", path]).decode("utf-8")
diff --git a/third_party/tsl/third_party/gpus/find_cuda_config.py b/third_party/tsl/third_party/gpus/find_cuda_config.py
index 68623bf671da71..c04dace79fe599 100644
--- a/third_party/tsl/third_party/gpus/find_cuda_config.py
+++ b/third_party/tsl/third_party/gpus/find_cuda_config.py
@@ -56,21 +56,15 @@
 tf_<library>_library_dir: ...
 """
 
+import glob
 import io
 import os
-import glob
 import platform
 import re
+import shutil
 import subprocess
 import sys
 
-# pylint: disable=g-import-not-at-top
-try:
-  from shutil import which
-except ImportError:
-  from distutils.spawn import find_executable as which
-# pylint: enable=g-import-not-at-top
-
 
 class ConfigError(Exception):
   pass
@@ -139,7 +133,7 @@ def _get_ld_config_paths():
   """Returns all directories from 'ldconfig -p'."""
   if not _is_linux():
     return []
-  ldconfig_path = which("ldconfig") or "/sbin/ldconfig"
+  ldconfig_path = shutil.which("ldconfig") or "/sbin/ldconfig"
   output = subprocess.check_output([ldconfig_path, "-p"])
   pattern = re.compile(".* => (.*)")
   result = set()