Fixed import for Python3