diff --git a/ci/_utils.sh b/ci/_utils.sh
index b4aae9cc7..9c0f4a847 100644
--- a/ci/_utils.sh
+++ b/ci/_utils.sh
@@ -234,7 +234,7 @@ start_message() {
     fi
     _rocm_path=`$REALPATH "$_rocm_path"`
     test -d "$_rocm_path" && echo "ROCm: $_rocm_path" || echo "ROCm path not found"
-    python --version
+    python3 --version
 }
 
 configure_omp_threads() {
@@ -266,6 +266,7 @@ pytest_run() {
     check_test_filter $_test_name_tag || return
     _start_ts=`date +%s`
     echo "Run [$_test_variant_tag] $@ at `time_elapsed $TEST_START_TS`"
-    pytest -v -rfEs `get_pytest_junitxml $_test_name_tag` $TEST_PYTEST_ARGS "$TEST_DIR/$@" || test_run_error "[$_test_variant_tag] $1"
+    python3 -m pytest -v -rfEs `get_pytest_junitxml $_test_name_tag` $TEST_PYTEST_ARGS "$TEST_DIR/$@"
+    test $? -eq 0 || test_run_error "[$_test_variant_tag] $1"
     echo "Done [$_test_variant_tag] $1 in `time_elapsed $_start_ts`"
 }
diff --git a/setup.py b/setup.py
index c66af39df..20ec983d3 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,10 @@
 
 import setuptools
 from setuptools.command.egg_info import egg_info
-from wheel.bdist_wheel import bdist_wheel
+try:
+    from setuptools.command.bdist_wheel import bdist_wheel
+except ImportError:
+    from wheel.bdist_wheel import bdist_wheel
 
 from build_tools.build_ext import CMakeExtension, get_build_ext
 from build_tools.te_version import te_version
diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py
index a80dcacb8..616a50f44 100644
--- a/transformer_engine/pytorch/setup.py
+++ b/transformer_engine/pytorch/setup.py
@@ -15,7 +15,10 @@
 import platform
 import urllib
 import setuptools
-from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+try:
+    from setuptools.command.bdist_wheel import bdist_wheel as _bdist_wheel
+except ImportError:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
 from packaging.version import parse
 
 try: