Fix import

resolve #5 @Simn