diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index af0b2f6..a1b8104 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -2115,7 +2115,8 @@
     status = psa_get_key_from_slot( key, &slot, PSA_KEY_USAGE_ENCRYPT, alg );
     if( status != PSA_SUCCESS )
         return( status );
-    if( ! PSA_KEY_TYPE_IS_KEYPAIR( slot->type ) )
+    if( ! ( PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->type ) ||
+            PSA_KEY_TYPE_IS_KEYPAIR( slot->type ) ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
 #if defined(MBEDTLS_RSA_C)