diff --git a/lib/psci/psci_common.c b/lib/psci/psci_common.c
index c893476..bfc09cc 100644
--- a/lib/psci/psci_common.c
+++ b/lib/psci/psci_common.c
@@ -451,8 +451,8 @@
  * enter. This function will be called after coordination of requested power
  * states has been done for each power level.
  *****************************************************************************/
-static void psci_set_target_local_pwr_states(unsigned int end_pwrlvl,
-					const psci_power_state_t *target_state)
+void psci_set_target_local_pwr_states(unsigned int end_pwrlvl,
+				      const psci_power_state_t *target_state)
 {
 	unsigned int parent_idx, lvl;
 	const plat_local_state_t *pd_state = target_state->pwr_domain_state;
@@ -474,7 +474,6 @@
 	}
 }
 
-
 /*******************************************************************************
  * PSCI helper function to get the parent nodes corresponding to a cpu_index.
  ******************************************************************************/
@@ -595,9 +594,6 @@
 		state_info->pwr_domain_state[lvl] = PSCI_LOCAL_STATE_RUN;
 
 	}
-
-	/* Update the target state in the power domain nodes */
-	psci_set_target_local_pwr_states(end_pwrlvl, state_info);
 }
 
 #if PSCI_OS_INIT_MODE
@@ -684,9 +680,6 @@
 		return rc;
 	}
 
-	/* Update the target state in the power domain nodes */
-	psci_set_target_local_pwr_states(end_pwrlvl, state_info);
-
 	return rc;
 }
 #endif
diff --git a/lib/psci/psci_off.c b/lib/psci/psci_off.c
index 9f36ac7..f83753f 100644
--- a/lib/psci/psci_off.c
+++ b/lib/psci/psci_off.c
@@ -104,6 +104,9 @@
 	 */
 	psci_do_state_coordination(end_pwrlvl, &state_info);
 
+	/* Update the target state in the power domain nodes */
+	psci_set_target_local_pwr_states(end_pwrlvl, &state_info);
+
 #if ENABLE_PSCI_STAT
 	/* Update the last cpu for each level till end_pwrlvl */
 	psci_stats_update_pwr_down(end_pwrlvl, &state_info);
diff --git a/lib/psci/psci_private.h b/lib/psci/psci_private.h
index b9987fe..04f93bd 100644
--- a/lib/psci/psci_private.h
+++ b/lib/psci/psci_private.h
@@ -298,6 +298,8 @@
 #endif
 void psci_get_target_local_pwr_states(unsigned int end_pwrlvl,
 				      psci_power_state_t *target_state);
+void psci_set_target_local_pwr_states(unsigned int end_pwrlvl,
+				      const psci_power_state_t *target_state);
 int psci_validate_entry_point(entry_point_info_t *ep,
 			uintptr_t entrypoint, u_register_t context_id);
 void psci_get_parent_pwr_domain_nodes(unsigned int cpu_idx,
diff --git a/lib/psci/psci_suspend.c b/lib/psci/psci_suspend.c
index 861b875..d93e60d 100644
--- a/lib/psci/psci_suspend.c
+++ b/lib/psci/psci_suspend.c
@@ -219,6 +219,19 @@
 	}
 #endif
 
+#if PSCI_OS_INIT_MODE
+	if (psci_plat_pm_ops->pwr_domain_validate_suspend != NULL) {
+		rc = psci_plat_pm_ops->pwr_domain_validate_suspend(state_info);
+		if (rc != PSCI_E_SUCCESS) {
+			skip_wfi = true;
+			goto exit;
+		}
+	}
+#endif
+
+	/* Update the target state in the power domain nodes */
+	psci_set_target_local_pwr_states(end_pwrlvl, state_info);
+
 #if ENABLE_PSCI_STAT
 	/* Update the last cpu for each level till end_pwrlvl */
 	psci_stats_update_pwr_down(end_pwrlvl, state_info);
@@ -234,15 +247,7 @@
 	 * program the power controller etc.
 	 */
 
-#if PSCI_OS_INIT_MODE
-	rc = psci_plat_pm_ops->pwr_domain_suspend(state_info);
-	if (rc != PSCI_E_SUCCESS) {
-		skip_wfi = true;
-		goto exit;
-	}
-#else
 	psci_plat_pm_ops->pwr_domain_suspend(state_info);
-#endif
 
 #if ENABLE_PSCI_STAT
 	plat_psci_stat_accounting_start(state_info);