diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5d8540c..98dad9d 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -6594,10 +6594,6 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_FFDH)
         case PSA_ALG_FFDH:
-            if (!PSA_KEY_TYPE_IS_DH_KEY_PAIR(psa_get_key_type(attributes))) {
-                return PSA_ERROR_INVALID_ARGUMENT;
-            }
-
             return mbedtls_psa_key_agreement_ffdh(attributes,
                                                   peer_key,
                                                   peer_key_length,
diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index d8c153a..6b07a7d 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -131,6 +131,10 @@
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    if (!PSA_KEY_TYPE_IS_DH_KEY_PAIR(psa_get_key_type(attributes))) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     mbedtls_mpi_init(&P); mbedtls_mpi_init(&G);
     mbedtls_mpi_init(&X); mbedtls_mpi_init(&GY);
     mbedtls_mpi_init(&K);