diff --git a/mol2vec/app/mol2vec.py b/mol2vec/app/mol2vec.py
index aab5d29..e1c7661 100644
--- a/mol2vec/app/mol2vec.py
+++ b/mol2vec/app/mol2vec.py
@@ -1,7 +1,11 @@
 import argparse
 import sys
+import six
 
-from mol2vec import features
+if six.PY2:
+    from .. import features
+else:
+    from mol2vec import features
 
 def do_corpus(args):
     print('Generating corpus file.')
@@ -166,4 +170,4 @@ def run():
 
 
 if __name__ == '__main__':
-    run()
\ No newline at end of file
+    run()
diff --git a/setup.py b/setup.py
index b575c1c..1c2a749 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,7 @@
       author_email='samo.turk@gmail.com, sabrina.jaeger@t-online.de, fulle@bio.mx',
       license='BSD 3-clause',
       packages=find_packages(),
-      install_requires=['numpy', 'gensim', 'tqdm', 'joblib', 'pandas', 'matplotlib', 'IPython', 'seaborn'],
+      install_requires=['numpy', 'gensim', 'tqdm', 'joblib', 'pandas', 'matplotlib', 'IPython', 'seaborn', 'six'],
       zip_safe=False,
       entry_points="""
       [console_scripts]