diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 44942911a0..376f6cf9a8 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1677,6 +1677,7 @@ int mbedtls_ssl_set_session(mbedtls_ssl_context *ssl, const mbedtls_ssl_session
             return MBEDTLS_ERR_SSL_BAD_INPUT_DATA;
         }
         ssl->handshake->ciphersuite_info = ciphersuite_info;
+        ssl->handshake->key_exchange_mode = MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ALL;
     }
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3 */
 
@@ -1686,7 +1687,6 @@ int mbedtls_ssl_set_session(mbedtls_ssl_context *ssl, const mbedtls_ssl_session
     }
 
     ssl->handshake->resume = 1;
-    ssl->handshake->key_exchange_mode = MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_ALL;
 
     return 0;
 }