static void unregister_nvdimm(void *nvdimm)
 {
+       struct cxl_nvdimm *cxl_nvd = nvdimm_provider_data(nvdimm);
+
        nvdimm_delete(nvdimm);
+       cxl_nvd->bridge = NULL;
 }
 
 static int cxl_nvdimm_probe(struct device *dev)
        }
 
        dev_set_drvdata(dev, nvdimm);
+       cxl_nvd->bridge = cxl_nvb;
        rc = devm_add_action_or_reset(dev, unregister_nvdimm, nvdimm);
 out:
        device_unlock(&cxl_nvb->dev);
        return cxl_nvb->nvdimm_bus != NULL;
 }
 
-static int cxl_nvdimm_release_driver(struct device *dev, void *data)
+static int cxl_nvdimm_release_driver(struct device *dev, void *cxl_nvb)
 {
+       struct cxl_nvdimm *cxl_nvd;
+
        if (!is_cxl_nvdimm(dev))
                return 0;
+
+       cxl_nvd = to_cxl_nvdimm(dev);
+       if (cxl_nvd->bridge != cxl_nvb)
+               return 0;
+
        device_release_driver(dev);
        return 0;
 }
 
-static void offline_nvdimm_bus(struct nvdimm_bus *nvdimm_bus)
+static void offline_nvdimm_bus(struct cxl_nvdimm_bridge *cxl_nvb,
+                              struct nvdimm_bus *nvdimm_bus)
 {
        if (!nvdimm_bus)
                return;
         * nvdimm_bus_unregister() rips the nvdimm objects out from
         * underneath them.
         */
-       bus_for_each_dev(&cxl_bus_type, NULL, NULL, cxl_nvdimm_release_driver);
+       bus_for_each_dev(&cxl_bus_type, NULL, cxl_nvb,
+                        cxl_nvdimm_release_driver);
        nvdimm_bus_unregister(nvdimm_bus);
 }
 
 
                dev_dbg(&cxl_nvb->dev, "rescan: %d\n", rc);
        }
-       offline_nvdimm_bus(victim_bus);
+       offline_nvdimm_bus(cxl_nvb, victim_bus);
 
        put_device(&cxl_nvb->dev);
 }