diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c5001f9..df0201b 100755
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1507,8 +1507,7 @@
 
     //TODO: check key policy
 
-    if ( !( ( key_type & PSA_KEY_TYPE_CATEGORY_MASK ) == PSA_KEY_TYPE_CATEGORY_SYMMETRIC
-            && PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) == cipher_info->block_size ) )
+    if ( ( key_type & PSA_KEY_TYPE_CATEGORY_MASK ) != PSA_KEY_TYPE_CATEGORY_SYMMETRIC )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
     if( alg == PSA_ALG_GCM )
@@ -1516,6 +1515,9 @@
         mbedtls_gcm_context gcm;
         tag_length = 16;
 
+        if( PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) != 16 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+
         //make sure we have place to hold the tag in the ciphertext buffer
         if( ciphertext_size < ( plaintext_length + tag_length ) )
             return( PSA_ERROR_INVALID_ARGUMENT );
@@ -1544,6 +1546,9 @@
         mbedtls_ccm_context ccm;
         tag_length = 16;
 
+        if( PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) != 16 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+
         if( nonce_length < 7 || nonce_length > 13 )
             return( PSA_ERROR_INVALID_ARGUMENT );