diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 2eb1d63e39..e319935a19 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -8184,7 +8184,7 @@ psa_status_t psa_generate_key_iop_complete(
 {
 #if defined(MBEDTLS_ECP_RESTARTABLE)
     psa_status_t status;
-    uint8_t key_data[MBEDTLS_ECP_MAX_BYTES] = { 0 };
+    uint8_t key_data[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(PSA_VENDOR_ECC_MAX_CURVE_BITS)] = { 0 };
     size_t key_len = 0;
 
     if (operation->id == 0 || operation->error_occurred) {