diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 1c066ce..e2cb06f 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7494,12 +7494,13 @@
 
     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);
 
     if (status == PSA_SUCCESS) {
-        operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
         if (operation->alg == PSA_ALG_JPAKE) {
             psa_jpake_computation_stage_t *computation_stage =
                 &operation->computation_stage.jpake;