diff --git a/IDE/STM32Cube/wolfssl_example.c b/IDE/STM32Cube/wolfssl_example.c
index 89c40ee6c..74f506a0a 100644
--- a/IDE/STM32Cube/wolfssl_example.c
+++ b/IDE/STM32Cube/wolfssl_example.c
@@ -977,7 +977,7 @@ static void client_thread(const void* args)
         osThreadSuspend(info->client.threadId);
 
         if (info->doShutdown)
-            info->client.done = 1;
+            info->server.done = 1;
     } while (!info->doShutdown);
 
     osThreadTerminate(info->client.threadId);
@@ -1252,7 +1252,7 @@ static void server_thread(const void* args)
         osThreadSuspend(info->server.threadId);
 
         if (info->doShutdown)
-            info->server.done = 1;
+            info->client.done = 1;
     } while (!info->doShutdown);
 
     osThreadTerminate(info->server.threadId);
diff --git a/examples/benchmark/tls_bench.c b/examples/benchmark/tls_bench.c
index e7d90fb2e..ba5ca2fb2 100644
--- a/examples/benchmark/tls_bench.c
+++ b/examples/benchmark/tls_bench.c
@@ -968,6 +968,19 @@ static int bench_tls_client(info_t* info)
     }
 #endif
 
+#ifndef NO_PSK
+    wolfSSL_CTX_set_psk_client_callback(cli_ctx, my_psk_client_cb);
+    #ifdef WOLFSSL_TLS13
+    #if !defined(WOLFSSL_PSK_TLS13_CB) && !defined(WOLFSSL_PSK_ONE_ID)
+    wolfSSL_CTX_set_psk_client_cs_callback(cli_ctx, my_psk_client_cs_cb);
+    #else
+    wolfSSL_CTX_set_psk_client_tls13_callback(cli_ctx, my_psk_client_tls13_cb);
+    #endif
+    #endif
+    wolfSSL_CTX_set_psk_callback_ctx(cli_ctx, (void*)info->cipher);
+#endif /* !NO_PSK */
+
+
     /* Allocate and initialize a packet sized buffer */
     writeBuf = (unsigned char*)XMALLOC(info->packetSize, NULL,
         DYNAMIC_TYPE_TMP_BUFFER);
@@ -1413,6 +1426,13 @@ static int bench_tls_server(info_t* info)
     }
 #endif
 
+#ifndef NO_PSK
+    wolfSSL_CTX_set_psk_server_callback(srv_ctx, my_psk_server_cb);
+    #ifdef WOLFSSL_TLS13
+    wolfSSL_CTX_set_psk_server_tls13_callback(srv_ctx, my_psk_server_tls13_cb);
+    #endif
+#endif /* !NO_PSK */
+
     /* Allocate read buffer */
     readBufSz = info->packetSize;
     readBuf = (unsigned char*)XMALLOC(readBufSz, NULL, DYNAMIC_TYPE_TMP_BUFFER);