return ret;
 }
 
+static void ad5064_bulk_reg_disable(void *data)
+{
+       struct ad5064_state *st = data;
+
+       regulator_bulk_disable(ad5064_num_vref(st), st->vref_reg);
+}
+
 static int ad5064_probe(struct device *dev, enum ad5064_type type,
                        const char *name, ad5064_write_func write)
 {
 
        st = iio_priv(indio_dev);
        mutex_init(&st->lock);
-       dev_set_drvdata(dev, indio_dev);
 
        st->chip_info = &ad5064_chip_info_tbl[type];
        st->dev = dev;
                ret = regulator_bulk_enable(ad5064_num_vref(st), st->vref_reg);
                if (ret)
                        return ret;
+
+               ret = devm_add_action_or_reset(dev, ad5064_bulk_reg_disable, st);
+               if (ret)
+                       return ret;
        }
 
        indio_dev->name = name;
                st->dac_cache[i] = midscale;
        }
 
-       ret = iio_device_register(indio_dev);
-       if (ret)
-               goto error_disable_reg;
-
-       return 0;
-
-error_disable_reg:
-       if (!st->use_internal_vref)
-               regulator_bulk_disable(ad5064_num_vref(st), st->vref_reg);
-
-       return ret;
-}
-
-static int ad5064_remove(struct device *dev)
-{
-       struct iio_dev *indio_dev = dev_get_drvdata(dev);
-       struct ad5064_state *st = iio_priv(indio_dev);
-
-       iio_device_unregister(indio_dev);
-
-       if (!st->use_internal_vref)
-               regulator_bulk_disable(ad5064_num_vref(st), st->vref_reg);
-
-       return 0;
+       return devm_iio_device_register(dev, indio_dev);
 }
 
 #if IS_ENABLED(CONFIG_SPI_MASTER)
                                ad5064_spi_write);
 }
 
-static int ad5064_spi_remove(struct spi_device *spi)
-{
-       return ad5064_remove(&spi->dev);
-}
-
 static const struct spi_device_id ad5064_spi_ids[] = {
        {"ad5024", ID_AD5024},
        {"ad5025", ID_AD5025},
                   .name = "ad5064",
        },
        .probe = ad5064_spi_probe,
-       .remove = ad5064_spi_remove,
        .id_table = ad5064_spi_ids,
 };
 
                                                ad5064_i2c_write);
 }
 
-static int ad5064_i2c_remove(struct i2c_client *i2c)
-{
-       return ad5064_remove(&i2c->dev);
-}
-
 static const struct i2c_device_id ad5064_i2c_ids[] = {
        {"ad5625", ID_AD5625 },
        {"ad5625r-1v25", ID_AD5625R_1V25 },
                   .name = "ad5064",
        },
        .probe = ad5064_i2c_probe,
-       .remove = ad5064_i2c_remove,
        .id_table = ad5064_i2c_ids,
 };