diff --git a/drivers/iommu/iommufd/main.c b/drivers/iommu/iommufd/main.c
index e71523cbd0de..e89c90784536 100644
--- a/drivers/iommu/iommufd/main.c
+++ b/drivers/iommu/iommufd/main.c
@@ -159,7 +159,8 @@ static struct iommufd_object *iommufd_object_remove(struct iommufd_ctx *ictx,
 	}
 
 	xas_store(&xas, NULL);
-	if (ictx->vfio_ioas == container_of(obj, struct iommufd_ioas, obj))
+	__xa_erase(&ictx->objects, obj->id);
+	if (ictx->vfio_ioas && &ictx->vfio_ioas->obj == obj)
 		ictx->vfio_ioas = NULL;
 
 out_xa: