diff --git a/src/common/defs.h b/src/common/defs.h
index 1de099f4b..21f1d1cd8 100644
--- a/src/common/defs.h
+++ b/src/common/defs.h
@@ -114,6 +114,8 @@ static inline int wpa_key_mgmt_sha256(int akm)
 {
 	return !!(akm & (WPA_KEY_MGMT_PSK_SHA256 |
 			 WPA_KEY_MGMT_IEEE8021X_SHA256 |
+			 WPA_KEY_MGMT_SAE |
+			 WPA_KEY_MGMT_FT_SAE |
 			 WPA_KEY_MGMT_OSEN |
 			 WPA_KEY_MGMT_IEEE8021X_SUITE_B |
 			 WPA_KEY_MGMT_FILS_SHA256 |
diff --git a/src/common/wpa_common.c b/src/common/wpa_common.c
index 13ae42272..f689fe83a 100644
--- a/src/common/wpa_common.c
+++ b/src/common/wpa_common.c
@@ -352,14 +352,14 @@ int wpa_pmk_to_ptk(const u8 *pmk, size_t pmk_len, const char *label,
 		return -1;
 #endif /* CONFIG_SUITEB192 || CONFIG_FILS */
 	} else if (wpa_key_mgmt_sha256(akmp) || akmp == WPA_KEY_MGMT_OWE) {
-#ifdef CONFIG_IEEE80211W
+#if defined(CONFIG_IEEE80211W) || defined(CONFIG_SAE)
 		wpa_printf(MSG_DEBUG, "WPA: PTK derivation using PRF(SHA256)");
 		if (sha256_prf(pmk, pmk_len, label, data, sizeof(data),
 			       tmp, ptk_len) < 0)
 			return -1;
-#else /* CONFIG_IEEE80211W */
+#else /* CONFIG_IEEE80211W or CONFIG_SAE */
 		return -1;
-#endif /* CONFIG_IEEE80211W */
+#endif /* CONFIG_IEEE80211W or CONFIG_SAE */
 #ifdef CONFIG_DPP
 	} else if (akmp == WPA_KEY_MGMT_DPP && pmk_len == 32) {
 		wpa_printf(MSG_DEBUG, "WPA: PTK derivation using PRF(SHA256)");