diff --git a/setup.py b/setup.py
index 459e6128a..7a870e265 100755
--- a/setup.py
+++ b/setup.py
@@ -5,15 +5,21 @@
 # Use of this source code is governed by a BSD-3-clause license that can be
 # found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
 
-import imp
 import os
 from setuptools import setup, find_packages
 
 # Get the coremltools version string
 coremltools_dir = os.path.join(os.path.dirname(__file__), "coremltools")
-version_module = imp.load_source(
-    "coremltools.version", os.path.join(coremltools_dir, "version.py")
-)
+
+try:
+    import importlib.util
+    spec = importlib.util.spec_from_file_location("coremltools.version", os.path.join(coremltools_dir, "version.py"))
+    version_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(version_module)
+except ImportError:  # Python 3.5 and below
+    import importlib.machinery
+    version_module = importlib.machinery.SourceFileLoader('coremltools.version', os.path.join(coremltools_dir, "version.py")).load_module()
+
 __version__ = version_module.__version__
 
 README = os.path.join(os.getcwd(), "README.md")