diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 9892660..68bea23 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1224,12 +1224,6 @@
 int mbedtls_ssl_tls13_finalize_write_client_hello(mbedtls_ssl_context *ssl)
 {
     ((void) ssl);
-#if defined(MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE)
-    mbedtls_ssl_handshake_set_state(
-        ssl, MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO);
-#else
-    mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_SERVER_HELLO);
-#endif /* MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE */
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
@@ -1239,6 +1233,10 @@
     const mbedtls_ssl_ciphersuite_t *ciphersuite_info;
 
     if (ssl->early_data_status == MBEDTLS_SSL_EARLY_DATA_STATUS_REJECTED) {
+#if defined(MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE)
+        mbedtls_ssl_handshake_set_state(
+            ssl, MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO);
+#endif
         MBEDTLS_SSL_DEBUG_MSG(
             1, ("Set hs psk for early data when writing the first psk"));