diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
index 0fc066f3c0b..90c25d55477 100644
--- a/tensorflow/tools/pip_package/setup.py
+++ b/tensorflow/tools/pip_package/setup.py
@@ -35,6 +35,7 @@ import re
 import sys
 
 from setuptools import Command
+from setuptools import Extension
 from setuptools import find_packages
 from setuptools import setup
 from setuptools.command.install import install as InstallCommandBase