Fix a import issue