diff --git a/arch/arm64/kvm/hisilicon/hisi_virt.c b/arch/arm64/kvm/hisilicon/hisi_virt.c
index d9029d612697c419af516ada58234d9cc763483c..48a8babb4daec5e337b06b5254b25bda35abb98a 100644
--- a/arch/arm64/kvm/hisilicon/hisi_virt.c
+++ b/arch/arm64/kvm/hisilicon/hisi_virt.c
@@ -258,10 +258,14 @@ int kvm_sched_affinity_vcpu_init(struct kvm_vcpu *vcpu)
 	if (!kvm_dvmbm_support)
 		return 0;
 
-	if (!zalloc_cpumask_var(&vcpu->arch.sched_cpus, GFP_ATOMIC) ||
-	    !zalloc_cpumask_var(&vcpu->arch.pre_sched_cpus, GFP_ATOMIC))
+	if (!zalloc_cpumask_var(&vcpu->arch.sched_cpus, GFP_ATOMIC))
 		return -ENOMEM;
 
+	if (!zalloc_cpumask_var(&vcpu->arch.pre_sched_cpus, GFP_ATOMIC)) {
+		free_cpumask_var(vcpu->arch.sched_cpus);
+		return -ENOMEM;
+	}
+
 	return 0;
 }