diff --git a/tools/keytools/keygen.c b/tools/keytools/keygen.c
index e1d644927..5f2d03648 100644
--- a/tools/keytools/keygen.c
+++ b/tools/keytools/keygen.c
@@ -1088,21 +1088,33 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
 
     if (exportPubKey) {
         if (saveAsDer) {
+            uint8_t*  pubDer;
+            size_t    pubDerSz;
+            int       pubOutLen;
+            const int WITH_ALG_SPKI = 1;
+
+            /* Size the buffer based on the ML DSA level */
+            switch (ml_dsa_level) {
+                case WC_ML_DSA_44:
+                    pubDerSz = ML_DSA_LEVEL2_PUB_KEY_DER_SIZE;
+                    break;
+                case WC_ML_DSA_65:
+                    pubDerSz = ML_DSA_LEVEL3_PUB_KEY_DER_SIZE;
+                    break;
+                case WC_ML_DSA_87:
+                    pubDerSz = ML_DSA_LEVEL5_PUB_KEY_DER_SIZE;
+                    break;
+                default:
+                    fprintf(stderr, "Error: Unsupported ML DSA level\n");
+                    exit(1);
+                    break;
+            }
+            pubDer = malloc(pubDerSz);
+
             /* Export public key in DER format */
-            uint8_t pubDer[
-            #if ML_DSA_LEVEL == 2
-                ML_DSA_LEVEL2_PUB_KEY_DER_SIZE
-            #elif ML_DSA_LEVEL == 3
-                ML_DSA_LEVEL3_PUB_KEY_DER_SIZE
-            #elif ML_DSA_LEVEL == 5
-                ML_DSA_LEVEL5_PUB_KEY_DER_SIZE
-            #endif
-            ];
-            int     pubOutLen;
 
-            const int WITH_ALG_SPKI = 1;
-            pubOutLen = wc_Dilithium_PublicKeyToDer(
-                &key, pubDer, sizeof(pubDer), WITH_ALG_SPKI);
+            pubOutLen = wc_Dilithium_PublicKeyToDer(&key, pubDer, pubDerSz,
+                                                    WITH_ALG_SPKI);
             if (pubOutLen < 0) {
                 fprintf(stderr, "Unable to export public key to DER, ret=%d\n",
                         pubOutLen);
@@ -1113,6 +1125,8 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
                 fprintf(stderr, "Unable to export public key to file\n");
                 exit(1);
             }
+
+            free(pubDer);
         }
         else {
             /* Export public key in raw format */