diff --git a/wolfcrypt/src/dilithium.c b/wolfcrypt/src/dilithium.c
index 625d6bf295..6a602a1043 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;
 }
 
@@ -7422,7 +7421,9 @@ int wc_dilithium_check_key(dilithium_key* key)
     }
 
     /* Dispose of allocated memory. */
-    XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    if (s1 != NULL) {
+        XFREE(s1, key->heap, DYNAMIC_TYPE_DILITHIUM);
+    }
 #else
     /* Validate parameter. */
     if (key == NULL) {
diff --git a/wolfcrypt/src/wc_kyber.c b/wolfcrypt/src/wc_kyber.c
index ffa37d84c1..a99cd6c1e9 100644
--- a/wolfcrypt/src/wc_kyber.c
+++ b/wolfcrypt/src/wc_kyber.c
@@ -286,7 +286,9 @@ int wc_KyberKey_MakeKeyWithRandom(KyberKey* key, const unsigned char* rand,
     }
 
     /* Free dynamic memory allocated in function. */
-    XFREE(a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
+    if (a != NULL) {
+        XFREE(a, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
+    }
 
     return ret;
 }
@@ -890,7 +892,9 @@ int wc_KyberKey_Decapsulate(KyberKey* key, unsigned char* ss,
 
 #ifndef USE_INTEL_SPEEDUP
     /* Dispose of dynamic memory allocated in function. */
-    XFREE(cmp, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
+    if (cmp != NULL) {
+        XFREE(cmp, key->heap, DYNAMIC_TYPE_TMP_BUFFER);
+    }
 #endif
 
     return ret;