1 files changed, 2 insertions, 2 deletions

diff --git a/drivers/infiniband/hw/hns/hns_roce_hem.c b/drivers/infiniband/hw/hns/hns_roce_hem.c
index 7487cf3d2c37..66f9f036ef94 100644
--- a/drivers/infiniband/hw/hns/hns_roce_hem.c
+++ b/drivers/infiniband/hw/hns/hns_roce_hem.c
@@ -1017,7 +1017,7 @@ void hns_roce_cleanup_hem_table(struct hns_roce_dev *hr_dev,
 
 void hns_roce_cleanup_hem(struct hns_roce_dev *hr_dev)
 {
-        if (hr_dev->caps.srqc_entry_sz)
+        if (hr_dev->caps.flags & HNS_ROCE_CAP_FLAG_SRQ)
                 hns_roce_cleanup_hem_table(hr_dev,
                                            &hr_dev->srq_table.table);
         hns_roce_cleanup_hem_table(hr_dev, &hr_dev->cq_table.table);
@@ -1027,7 +1027,7 @@ void hns_roce_cleanup_hem(struct hns_roce_dev *hr_dev)
         if (hr_dev->caps.cqc_timer_entry_sz)
                 hns_roce_cleanup_hem_table(hr_dev,
                                            &hr_dev->cqc_timer_table);
-        if (hr_dev->caps.sccc_sz)
+        if (hr_dev->caps.flags & HNS_ROCE_CAP_FLAG_QP_FLOW_CTRL)
                 hns_roce_cleanup_hem_table(hr_dev,
                                            &hr_dev->qp_table.sccc_table);
         if (hr_dev->caps.trrl_entry_sz)