diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index a3d33a3..76f0f18 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1236,10 +1236,6 @@
     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"));
 
@@ -1294,6 +1290,15 @@
             return ret;
         }
 
+#if defined(MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE)
+        mbedtls_ssl_handshake_set_state(
+            ssl, MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO);
+#else
+        MBEDTLS_SSL_DEBUG_MSG(
+            1, ("Switch to early data keys for outbound traffic"));
+        mbedtls_ssl_set_outbound_transform(
+            ssl, ssl->handshake->transform_earlydata);
+#endif
     }
 #endif /* MBEDTLS_SSL_EARLY_DATA */
     return 0;
@@ -3067,19 +3072,19 @@
             }
             break;
 
+#if defined(MBEDTLS_SSL_EARLY_DATA)
         case MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO:
             ret = mbedtls_ssl_tls13_write_change_cipher_spec(ssl);
             if (ret == 0) {
                 mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_SERVER_HELLO);
 
-#if defined(MBEDTLS_SSL_EARLY_DATA)
                 MBEDTLS_SSL_DEBUG_MSG(
                     1, ("Switch to early data keys for outbound traffic"));
                 mbedtls_ssl_set_outbound_transform(
                     ssl, ssl->handshake->transform_earlydata);
-#endif
             }
             break;
+#endif /* MBEDTLS_SSL_EARLY_DATA */
 #endif /* MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE */
 
 #if defined(MBEDTLS_SSL_SESSION_TICKETS)