diff --git a/lib/test_hmm.c b/lib/test_hmm.c
index 79fe7d233df1..213504915737 100644
--- a/lib/test_hmm.c
+++ b/lib/test_hmm.c
@@ -1738,6 +1738,13 @@ static const struct dev_pagemap_ops dmirror_devmem_ops = {
 	.folio_split	= dmirror_devmem_folio_split,
 };
 
+static void dmirror_device_release(struct device *dev)
+{
+	struct dmirror_device *mdevice = container_of(dev, struct dmirror_device, device);
+
+	dmirror_device_remove_chunks(mdevice);
+}
+
 static int dmirror_device_init(struct dmirror_device *mdevice, int id)
 {
 	dev_t dev;
@@ -1749,6 +1756,8 @@ static int dmirror_device_init(struct dmirror_device *mdevice, int id)
 
 	cdev_init(&mdevice->cdevice, &dmirror_fops);
 	mdevice->cdevice.owner = THIS_MODULE;
+	mdevice->device.release = dmirror_device_release;
+
 	device_initialize(&mdevice->device);
 	mdevice->device.devt = dev;
 
@@ -1756,12 +1765,16 @@ static int dmirror_device_init(struct dmirror_device *mdevice, int id)
 	if (ret)
 		goto put_device;
 
+	/* Build a list of free ZONE_DEVICE struct pages */
+	ret = dmirror_allocate_chunk(mdevice, NULL, false);
+	if (ret)
+		goto put_device;
+
 	ret = cdev_device_add(&mdevice->cdevice, &mdevice->device);
 	if (ret)
 		goto put_device;
 
-	/* Build a list of free ZONE_DEVICE struct pages */
-	return dmirror_allocate_chunk(mdevice, NULL, false);
+	return 0;
 
 put_device:
 	put_device(&mdevice->device);
@@ -1770,7 +1783,6 @@ static int dmirror_device_init(struct dmirror_device *mdevice, int id)
 
 static void dmirror_device_remove(struct dmirror_device *mdevice)
 {
-	dmirror_device_remove_chunks(mdevice);
 	cdev_device_del(&mdevice->cdevice, &mdevice->device);
 	put_device(&mdevice->device);
 }