diff --git a/wolfcrypt/src/dilithium.c b/wolfcrypt/src/dilithium.c
index 625d6bf29..bd5b10272 100644
--- a/wolfcrypt/src/dilithium.c
+++ b/wolfcrypt/src/dilithium.c
@@ -6952,10 +6952,9 @@ int wc_dilithium_init_ex(dilithium_key* key, void* heap, int devId)
         key->idLen = 0;
         key->labelLen = 0;
     #endif
+        key->heap = heap;
     }
 
-    key->heap = heap;
-
     return ret;
 }