diff --git a/library/lms.c b/library/lms.c
index ccbcd8b..35ca890 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -550,17 +550,24 @@
 
     if( ctx->have_private_key )
     {
-        for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM(ctx->params.type); idx++ )
+        if( ctx->ots_private_keys != NULL )
         {
-            mbedtls_lmots_private_free( &ctx->ots_private_keys[idx] );
-            mbedtls_lmots_public_free( &ctx->ots_public_keys[idx] );
+            for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM(ctx->params.type); idx++ )
+            {
+                mbedtls_lmots_private_free( &ctx->ots_private_keys[idx] );
+            }
         }
 
-        if( ctx->ots_private_keys != NULL )
-            mbedtls_free( ctx->ots_private_keys );
-
         if( ctx->ots_public_keys != NULL )
-            mbedtls_free( ctx->ots_public_keys );
+        {
+        for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM(ctx->params.type); idx++ )
+            {
+                mbedtls_lmots_public_free( &ctx->ots_public_keys[idx] );
+            }
+        }
+
+        mbedtls_free( ctx->ots_private_keys );
+        mbedtls_free( ctx->ots_public_keys );
     }
 
     mbedtls_platform_zeroize( ctx, sizeof( *ctx ) );
@@ -594,6 +601,7 @@
 
     ctx->params.type = type;
     ctx->params.otstype = otstype;
+    ctx->have_private_key = 1;
 
     ret = f_rng( p_rng,
                  ctx->params.I_key_identifier,
@@ -619,22 +627,10 @@
                                            sizeof( *ctx->ots_public_keys ) );
     if( ctx->ots_public_keys == NULL )
     {
-        /* Free just the ots private keys (since they've been allocated at this
-         * point) so that we can pass the context to lms_private_free (which
-         * will not try to free the private keys since have_private_key is not
-         * set.
-         */
-        mbedtls_free(ctx->ots_private_keys);
-        ctx->ots_private_keys = NULL;
         ret = MBEDTLS_ERR_LMS_ALLOC_FAILED;
         goto exit;
     }
 
-    /* Now that all the allocation has succeeded we set have_private_key, since
-     * that causes lms_private_free to free the ots keys.
-     */
-    ctx->have_private_key = 1;
-
     for( idx = 0; idx < MERKLE_TREE_LEAF_NODE_AM(ctx->params.type); idx++ )
     {
         mbedtls_lmots_private_init( &ctx->ots_private_keys[idx] );