diff --git a/docs/source/conf.py b/docs/source/conf.py
index 5e0b468c..4837e7ca 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -27,9 +27,9 @@
 author = 'Yahoo Inc'
 
 # The short X.Y version
-version = '1.4.0'
+version = '1.4.1'
 # The full version, including alpha/beta/rc tags
-release = '1.4.0'
+release = '1.4.1'
 
 
 # -- General configuration ---------------------------------------------------
diff --git a/setup.py b/setup.py
index 2a04d70d..5f4ff4e5 100644
--- a/setup.py
+++ b/setup.py
@@ -1,16 +1,21 @@
 from setuptools import setup
 
+with open('README.md') as f:
+  long_description = f.read()
+
 setup(
-  name = 'tensorflowonspark',
-  packages = ['tensorflowonspark'],
-  version = '1.4.0',
-  description = 'Deep learning with TensorFlow on Apache Spark clusters',
-  author = 'Yahoo, Inc.',
-  url = 'https://github.com/yahoo/TensorFlowOnSpark',
-  keywords = ['tensorflowonspark', 'tensorflow', 'spark', 'machine learning', 'yahoo'],
-  install_requires = ['tensorflow'],
-  license = 'Apache 2.0',
-  classifiers = [
+  name='tensorflowonspark',
+  packages=['tensorflowonspark'],
+  version='1.4.1',
+  description='Deep learning with TensorFlow on Apache Spark clusters',
+  long_description=long_description,
+  long_description_content_type='text/markdown',
+  author='Yahoo, Inc.',
+  url='https://github.com/yahoo/TensorFlowOnSpark',
+  keywords=['tensorflowonspark', 'tensorflow', 'spark', 'machine learning', 'yahoo'],
+  install_requires=['tensorflow'],
+  license='Apache 2.0',
+  classifiers=[
     'Development Status :: 4 - Beta',
     'Intended Audience :: Developers',
     'Intended Audience :: Science/Research',
diff --git a/tensorflowonspark/__init__.py b/tensorflowonspark/__init__.py
index 33b5f8e6..257a5178 100644
--- a/tensorflowonspark/__init__.py
+++ b/tensorflowonspark/__init__.py
@@ -2,4 +2,4 @@
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s (%(threadName)s-%(process)d) %(message)s")
 
-__version__ = "1.4.0"
+__version__ = "1.4.1"
diff --git a/tensorflowonspark/util.py b/tensorflowonspark/util.py
index dc167226..5ebaa83b 100644
--- a/tensorflowonspark/util.py
+++ b/tensorflowonspark/util.py
@@ -10,8 +10,11 @@
 import logging
 import os
 import socket
+import subprocess
 import errno
 from socket import error as socket_error
+from . import gpu_info
+
 
 def single_node_env(num_gpus=1):
   """Setup environment variables for Hadoop compatibility and GPU allocation"""
@@ -34,12 +37,13 @@ def single_node_env(num_gpus=1):
     logging.info("Using CPU")
     os.environ['CUDA_VISIBLE_DEVICES'] = ''
 
+
 def get_ip_address():
   """Simple utility to get host IP address."""
   try:
     s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
     s.connect(("8.8.8.8", 80))
-    ip_address =  s.getsockname()[0]
+    ip_address = s.getsockname()[0]
   except socket_error as sockerr:
     if sockerr.errno != errno.ENETUNREACH:
       raise sockerr