diff --git a/drivers/platform/x86/lenovo-yogabook-wmi.c b/drivers/platform/x86/lenovo-yogabook-wmi.c
index c5d103099ea9..4bcbf0f84732 100644
--- a/drivers/platform/x86/lenovo-yogabook-wmi.c
+++ b/drivers/platform/x86/lenovo-yogabook-wmi.c
@@ -200,16 +200,11 @@ static struct led_lookup_data yogabook_pen_led = {
 	.con_id = "pen-icon-led",
 };
 
-static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
+static int yogabook_probe(struct device *dev, struct yogabook_wmi *data,
+			  const char *kbd_bl_led_name)
 {
-	struct device *dev = &wdev->dev;
-	struct yogabook_wmi *data;
 	int r;
 
-	data = devm_kzalloc(dev, sizeof(struct yogabook_wmi), GFP_KERNEL);
-	if (data == NULL)
-		return -ENOMEM;
-
 	dev_set_drvdata(dev, data);
 
 	data->dev = dev;
@@ -218,31 +213,6 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 	set_bit(YB_DIGITIZER_IS_ON, &data->flags);
 	INIT_WORK(&data->work, yogabook_wmi_work);
 
-	data->kbd_adev = acpi_dev_get_first_match_dev("GDIX1001", NULL, -1);
-	if (!data->kbd_adev) {
-		dev_err(dev, "Cannot find the touchpad device in ACPI tables\n");
-		return -ENODEV;
-	}
-
-	data->dig_adev = acpi_dev_get_first_match_dev("WCOM0019", NULL, -1);
-	if (!data->dig_adev) {
-		dev_err(dev, "Cannot find the digitizer device in ACPI tables\n");
-		r = -ENODEV;
-		goto error_put_devs;
-	}
-
-	data->kbd_dev = get_device(acpi_get_first_physical_node(data->kbd_adev));
-	if (!data->kbd_dev || !data->kbd_dev->driver) {
-		r = -EPROBE_DEFER;
-		goto error_put_devs;
-	}
-
-	data->dig_dev = get_device(acpi_get_first_physical_node(data->dig_adev));
-	if (!data->dig_dev || !data->dig_dev->driver) {
-		r = -EPROBE_DEFER;
-		goto error_put_devs;
-	}
-
 	yogabook_pen_led.dev_id = dev_name(dev);
 	led_add_lookup(&yogabook_pen_led);
 	data->pen_led = devm_led_get(dev, "pen-icon-led");
@@ -251,7 +221,7 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 	if (IS_ERR(data->pen_led)) {
 		r = PTR_ERR(data->pen_led);
 		dev_err_probe(dev, r, "Getting pen icon LED\n");
-		goto error_put_devs;
+		return r;
 	}
 
 	yogabook_wmi_gpios.dev_id = dev_name(dev);
@@ -262,13 +232,13 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 	if (IS_ERR(data->backside_hall_gpio)) {
 		r = PTR_ERR(data->backside_hall_gpio);
 		dev_err_probe(dev, r, "Getting backside_hall_sw GPIO\n");
-		goto error_put_devs;
+		return r;
 	}
 
 	r = gpiod_to_irq(data->backside_hall_gpio);
 	if (r < 0) {
 		dev_err_probe(dev, r, "Getting backside_hall_sw IRQ\n");
-		goto error_put_devs;
+		return r;
 	}
 	data->backside_hall_irq = r;
 
@@ -280,12 +250,12 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 			"backside_hall_sw", data);
 	if (r) {
 		dev_err_probe(dev, r, "Requesting backside_hall_sw IRQ\n");
-		goto error_put_devs;
+		return r;
 	}
 
 	schedule_work(&data->work);
 
-	data->kbd_bl_led.name = "ybwmi::kbd_backlight";
+	data->kbd_bl_led.name = kbd_bl_led_name;
 	data->kbd_bl_led.brightness_set_blocking = kbd_brightness_set;
 	data->kbd_bl_led.brightness_get = kbd_brightness_get;
 	data->kbd_bl_led.max_brightness = 255;
@@ -301,6 +271,50 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 error_free_irq:
 	free_irq(data->backside_hall_irq, data);
 	cancel_work_sync(&data->work);
+	return r;
+}
+
+static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
+{
+	struct device *dev = &wdev->dev;
+	struct yogabook_wmi *data;
+	int r;
+
+	data = devm_kzalloc(dev, sizeof(*data), GFP_KERNEL);
+	if (data == NULL)
+		return -ENOMEM;
+
+	data->kbd_adev = acpi_dev_get_first_match_dev("GDIX1001", NULL, -1);
+	if (!data->kbd_adev) {
+		dev_err(dev, "Cannot find the touchpad device in ACPI tables\n");
+		return -ENODEV;
+	}
+
+	data->dig_adev = acpi_dev_get_first_match_dev("WCOM0019", NULL, -1);
+	if (!data->dig_adev) {
+		dev_err(dev, "Cannot find the digitizer device in ACPI tables\n");
+		r = -ENODEV;
+		goto error_put_devs;
+	}
+
+	data->kbd_dev = get_device(acpi_get_first_physical_node(data->kbd_adev));
+	if (!data->kbd_dev || !data->kbd_dev->driver) {
+		r = -EPROBE_DEFER;
+		goto error_put_devs;
+	}
+
+	data->dig_dev = get_device(acpi_get_first_physical_node(data->dig_adev));
+	if (!data->dig_dev || !data->dig_dev->driver) {
+		r = -EPROBE_DEFER;
+		goto error_put_devs;
+	}
+
+	r = yogabook_probe(dev, data, "ybwmi::kbd_backlight");
+	if (r)
+		goto error_put_devs;
+
+	return 0;
+
 error_put_devs:
 	put_device(data->dig_dev);
 	put_device(data->kbd_dev);
@@ -309,9 +323,8 @@ static int yogabook_wmi_probe(struct wmi_device *wdev, const void *context)
 	return r;
 }
 
-static void yogabook_wmi_remove(struct wmi_device *wdev)
+static void yogabook_remove(struct yogabook_wmi *data)
 {
-	struct yogabook_wmi *data = dev_get_drvdata(&wdev->dev);
 	int r = 0;
 
 	free_irq(data->backside_hall_irq, data);
@@ -325,6 +338,13 @@ static void yogabook_wmi_remove(struct wmi_device *wdev)
 
 	if (r)
 		dev_warn(data->dev, "Reprobe of devices failed\n");
+}
+
+static void yogabook_wmi_remove(struct wmi_device *wdev)
+{
+	struct yogabook_wmi *data = dev_get_drvdata(&wdev->dev);
+
+	yogabook_remove(data);
 
 	put_device(data->dig_dev);
 	put_device(data->kbd_dev);
