diff --git a/drivers/crypto/hisilicon/migration/acc_vf_migration.c b/drivers/crypto/hisilicon/migration/acc_vf_migration.c
index fdfbe4590253d507bddb770d40e2a0ea227c1dd7..cdaa55aa8716650fff1feb19c8c3153e368ecb65 100644
--- a/drivers/crypto/hisilicon/migration/acc_vf_migration.c
+++ b/drivers/crypto/hisilicon/migration/acc_vf_migration.c
@@ -1779,7 +1779,9 @@ static int __init acc_vf_module_init(void)
 
 static void __exit acc_vf_module_exit(void)
 {
-	vfio_pci_unregister_vendor_driver(&acc_vf_device_ops_node);
+	vfio_pci_unregister_vendor_driver(&sec_vf_mig_ops);
+	vfio_pci_unregister_vendor_driver(&hpre_vf_mig_ops);
+	vfio_pci_unregister_vendor_driver(&zip_vf_mig_ops);
 };
 module_init(acc_vf_module_init);
 module_exit(acc_vf_module_exit);
diff --git a/drivers/vfio/pci/vfio_pci.c b/drivers/vfio/pci/vfio_pci.c
index 2382222b096051c4053d86c0ab3dc0104f25305f..c04b34247067454c53df42d800e9bf0efdb570fb 100644
--- a/drivers/vfio/pci/vfio_pci.c
+++ b/drivers/vfio/pci/vfio_pci.c
@@ -2608,7 +2608,8 @@ int __vfio_pci_register_vendor_driver(struct vfio_pci_vendor_driver_ops *ops)
 
 	/* Check for duplicates */
 	list_for_each_entry(tmp, &vfio_pci.vendor_drivers_list, next) {
-		if (tmp->ops->device_ops == ops->device_ops) {
+		if (tmp->ops->vendor == ops->vendor &&
+		    tmp->ops->vendor == ops->device) {
 			mutex_unlock(&vfio_pci.vendor_drivers_lock);
 			kfree(driver);
 			return -EINVAL;
@@ -2626,14 +2627,15 @@ int __vfio_pci_register_vendor_driver(struct vfio_pci_vendor_driver_ops *ops)
 }
 EXPORT_SYMBOL_GPL(__vfio_pci_register_vendor_driver);
 
-void vfio_pci_unregister_vendor_driver(struct vfio_device_ops *device_ops)
+void vfio_pci_unregister_vendor_driver(struct vfio_pci_vendor_driver_ops *ops)
 {
 	struct vfio_pci_vendor_driver *driver, *tmp;
 
 	mutex_lock(&vfio_pci.vendor_drivers_lock);
 	list_for_each_entry_safe(driver, tmp,
 				 &vfio_pci.vendor_drivers_list, next) {
-		if (driver->ops->device_ops == device_ops) {
+		if (driver->ops->vendor == ops->vendor &&
+		    driver->ops->device == ops->device) {
 			list_del(&driver->next);
 			mutex_unlock(&vfio_pci.vendor_drivers_lock);
 			kfree(driver);
diff --git a/include/linux/vfio.h b/include/linux/vfio.h
index e8841446a2eba81d3b6b73f9ca8536b807645a2d..f516784e0356f866ed36110d433df58934184edc 100644
--- a/include/linux/vfio.h
+++ b/include/linux/vfio.h
@@ -261,7 +261,7 @@ struct vfio_pci_vendor_driver_ops {
 	struct vfio_device_ops *device_ops;
 };
 int __vfio_pci_register_vendor_driver(struct vfio_pci_vendor_driver_ops *ops);
-void vfio_pci_unregister_vendor_driver(struct vfio_device_ops *device_ops);
+void vfio_pci_unregister_vendor_driver(struct vfio_pci_vendor_driver_ops *ops);
 
 #define vfio_pci_register_vendor_driver(__name, __probe, __remove,	\
 					__device_ops)			\