This makes import work correctly after installing with pip install . or setup.py install.
pip install .
setup.py install