diff --git a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
index 82e873680e..4500196ef0 100644
--- a/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
+++ b/tf-psa-crypto/drivers/builtin/src/psa_crypto_ecp.c
@@ -625,14 +625,6 @@ psa_status_t mbedtls_psa_generate_key_iop_complete(
 {
     *key_len = 0;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
-                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
-
-    if (status) {
-        return mbedtls_to_psa_error(status);
-    }
-
-    operation->num_ops = 1;
 
     *key_len = PSA_BITS_TO_BYTES(operation->ecp.grp.nbits);
 
@@ -640,6 +632,15 @@ psa_status_t mbedtls_psa_generate_key_iop_complete(
         return PSA_ERROR_BUFFER_TOO_SMALL;
     }
 
+    status = mbedtls_ecp_gen_privkey(&operation->ecp.grp, &operation->ecp.d,
+                                     mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE);
+
+    if (status != 0) {
+        return mbedtls_to_psa_error(status);
+    }
+
+    operation->num_ops = 1;
+
     mbedtls_mpi_write_binary(&operation->ecp.d, key_output, key_output_size);
 
     return mbedtls_to_psa_error(status);