}
 }
 
+static void bma400_regulators_disable(void *data_ptr)
+{
+       struct bma400_data *data = data_ptr;
+
+       regulator_bulk_disable(ARRAY_SIZE(data->regulators), data->regulators);
+}
+
+static void bma400_power_disable(void *data_ptr)
+{
+       struct bma400_data *data = data_ptr;
+       int ret;
+
+       mutex_lock(&data->mutex);
+       ret = bma400_set_power_mode(data, POWER_MODE_SLEEP);
+       mutex_unlock(&data->mutex);
+       if (ret)
+               dev_warn(data->dev, "Failed to put device into sleep mode (%pe)\n",
+                        ERR_PTR(ret));
+}
+
 static int bma400_init(struct bma400_data *data)
 {
        unsigned int val;
        ret = regmap_read(data->regmap, BMA400_CHIP_ID_REG, &val);
        if (ret) {
                dev_err(data->dev, "Failed to read chip id register\n");
-               goto out;
+               return ret;
        }
 
        if (val != BMA400_ID_REG_VAL) {
                dev_err(data->dev, "Chip ID mismatch\n");
-               ret = -ENODEV;
-               goto out;
+               return -ENODEV;
        }
 
        data->regulators[BMA400_VDD_REGULATOR].supply = "vdd";
                                "Failed to get regulators: %d\n",
                                ret);
 
-               goto out;
+               return ret;
        }
        ret = regulator_bulk_enable(ARRAY_SIZE(data->regulators),
                                    data->regulators);
        if (ret) {
                dev_err(data->dev, "Failed to enable regulators: %d\n",
                        ret);
-               goto out;
+               return ret;
        }
 
+       ret = devm_add_action_or_reset(data->dev, bma400_regulators_disable, data);
+       if (ret)
+               return ret;
+
        ret = bma400_get_power_mode(data);
        if (ret) {
                dev_err(data->dev, "Failed to get the initial power-mode\n");
-               goto err_reg_disable;
+               return ret;
        }
 
        if (data->power_mode != POWER_MODE_NORMAL) {
                ret = bma400_set_power_mode(data, POWER_MODE_NORMAL);
                if (ret) {
                        dev_err(data->dev, "Failed to wake up the device\n");
-                       goto err_reg_disable;
+                       return ret;
                }
                /*
                 * TODO: The datasheet waits 1500us here in the example, but
                usleep_range(1500, 2000);
        }
 
+       ret = devm_add_action_or_reset(data->dev, bma400_power_disable, data);
+       if (ret)
+               return ret;
+
        bma400_init_tables();
 
        ret = bma400_get_accel_output_data_rate(data);
        if (ret)
-               goto err_reg_disable;
+               return ret;
 
        ret = bma400_get_accel_oversampling_ratio(data);
        if (ret)
-               goto err_reg_disable;
+               return ret;
 
        ret = bma400_get_accel_scale(data);
        if (ret)
-               goto err_reg_disable;
+               return ret;
 
        /*
         * Once the interrupt engine is supported we might use the
         * channel.
         */
        return regmap_write(data->regmap, BMA400_ACC_CONFIG2_REG, 0x00);
-
-err_reg_disable:
-       regulator_bulk_disable(ARRAY_SIZE(data->regulators),
-                              data->regulators);
-out:
-       return ret;
 }
 
 static int bma400_read_raw(struct iio_dev *indio_dev,
        indio_dev->num_channels = ARRAY_SIZE(bma400_channels);
        indio_dev->modes = INDIO_DIRECT_MODE;
 
-       dev_set_drvdata(dev, indio_dev);
-
-       return iio_device_register(indio_dev);
+       return devm_iio_device_register(dev, indio_dev);
 }
 EXPORT_SYMBOL_NS(bma400_probe, IIO_BMA400);
 
-void bma400_remove(struct device *dev)
-{
-       struct iio_dev *indio_dev = dev_get_drvdata(dev);
-       struct bma400_data *data = iio_priv(indio_dev);
-       int ret;
-
-       mutex_lock(&data->mutex);
-       ret = bma400_set_power_mode(data, POWER_MODE_SLEEP);
-       mutex_unlock(&data->mutex);
-
-       if (ret)
-               dev_warn(dev, "Failed to put device into sleep mode (%pe)\n", ERR_PTR(ret));
-
-       regulator_bulk_disable(ARRAY_SIZE(data->regulators),
-                              data->regulators);
-
-       iio_device_unregister(indio_dev);
-}
-EXPORT_SYMBOL_NS(bma400_remove, IIO_BMA400);
-
 MODULE_AUTHOR("Dan Robertson <dan@dlrobertson.com>");
 MODULE_DESCRIPTION("Bosch BMA400 triaxial acceleration sensor core");
 MODULE_LICENSE("GPL");