Fixed import error in Python 3

Added conditional to detect Python 3.