Fixed import for python3