diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 1438124..810db24 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1932,11 +1932,6 @@
     if( ssl->handshake == NULL || ssl->conf == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
-    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
-        psa_role = PSA_PAKE_ROLE_SERVER;
-    else
-        psa_role = PSA_PAKE_ROLE_CLIENT;
-
     if( mbedtls_svc_key_id_is_null( pwd ) )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
     ssl->handshake->psa_pake_password = pwd;
@@ -1952,6 +1947,11 @@
     if( status != PSA_SUCCESS )
         goto error;
 
+    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
+        psa_role = PSA_PAKE_ROLE_SERVER;
+    else
+        psa_role = PSA_PAKE_ROLE_CLIENT;
+
     status = psa_pake_set_role( &ssl->handshake->psa_pake_ctx, psa_role );
     if( status != PSA_SUCCESS )
         goto error;
@@ -1979,15 +1979,15 @@
     if( ssl->handshake == NULL || ssl->conf == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
+    /* Empty password is not valid  */
+    if( ( pw == NULL) || ( pw_len == 0 ) )
+        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
+
     if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
         role = MBEDTLS_ECJPAKE_SERVER;
     else
         role = MBEDTLS_ECJPAKE_CLIENT;
 
-    /* Empty password is not valid  */
-    if( ( pw == NULL) || ( pw_len == 0 ) )
-        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
-
     return( mbedtls_ecjpake_setup( &ssl->handshake->ecjpake_ctx,
                                    role,
                                    MBEDTLS_MD_SHA256,