Revert "Remove useless import in __init__"

This reverts commit ab2732e8f26c9b4a91fe7ec4fdff0e57c95e0367.