diff --git a/library/lms.c b/library/lms.c
index a4235ad..fba5d88 100644
--- a/library/lms.c
+++ b/library/lms.c
@@ -235,11 +235,6 @@
     mbedtls_lms_algorithm_type_t type;
     mbedtls_lmots_algorithm_type_t otstype;
 
-    if( key_size != MBEDTLS_LMS_PUBLIC_KEY_LEN(ctx->params.type) )
-    {
-        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
-    }
-
     type = mbedtls_lms_network_bytes_to_unsigned_int( MBEDTLS_LMS_TYPE_LEN,
             key + PUBLIC_KEY_TYPE_OFFSET );
     if( type != MBEDTLS_LMS_SHA256_M32_H10 )
@@ -248,6 +243,11 @@
     }
     ctx->params.type = type;
 
+    if( key_size != MBEDTLS_LMS_PUBLIC_KEY_LEN(ctx->params.type) )
+    {
+        return( MBEDTLS_ERR_LMS_BAD_INPUT_DATA );
+    }
+
     otstype = mbedtls_lms_network_bytes_to_unsigned_int( MBEDTLS_LMOTS_TYPE_LEN,
             key + PUBLIC_KEY_OTSTYPE_OFFSET );
     if( otstype != MBEDTLS_LMOTS_SHA256_N32_W8 )