diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 4406fcc3b..8752bffe5 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7503,8 +7503,6 @@ static psa_status_t psa_pake_complete_inputs(
 
     status = psa_driver_wrapper_pake_setup(operation, &inputs);
 
-    operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
-
     /* Driver is responsible for creating its own copy of the password. */
     mbedtls_platform_zeroize(inputs.password, inputs.password_len);
     mbedtls_free(inputs.password);
@@ -7512,6 +7510,7 @@ static psa_status_t psa_pake_complete_inputs(
     if (status == PSA_SUCCESS) {
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
         if (operation->alg == PSA_ALG_JPAKE) {
+            operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
             psa_jpake_computation_stage_t *computation_stage =
                 &operation->computation_stage.jpake;
             computation_stage->state = PSA_PAKE_STATE_READY;
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 062d0bb68..63d08303d 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -257,9 +257,11 @@ psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
     { status = PSA_ERROR_NOT_SUPPORTED; }
 
 error:
-    /* When driver fails with PSA_ERROR_NOT_SUPPORTED the built-in implementation is executed (if available)
-       and it will reallocate the password leading to the memory leak.
-       Call abort explicitly to clean up allocated memory for password on failure. */
+    /* In case of failure of the setup of a multipart operation, the PSA driver interface
+     * specifies that the core does not call any other driver entry point thus does not
+     * call mbedtls_psa_pake_abort(). Therefore call it here to do the needed clean
+     * up like freeing the memory that may have been allocated to store the password.
+     */
     mbedtls_psa_pake_abort(operation);
     return status;
 }