diff --git a/include/psa/crypto_builtin_key_derivation.h b/include/psa/crypto_builtin_key_derivation.h
index d2cf4df97..5d01f6c58 100644
--- a/include/psa/crypto_builtin_key_derivation.h
+++ b/include/psa/crypto_builtin_key_derivation.h
@@ -116,7 +116,7 @@ typedef enum {
 
 typedef struct {
     psa_pbkdf2_key_derivation_state_t MBEDTLS_PRIVATE(state);
-    uint64_t MBEDTLS_PRIVATE(input_cost);
+    size_t MBEDTLS_PRIVATE(input_cost);
     uint8_t *MBEDTLS_PRIVATE(salt);
     size_t MBEDTLS_PRIVATE(salt_length);
     uint8_t *MBEDTLS_PRIVATE(password);
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 59169d504..af4ab6515 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -6423,11 +6423,11 @@ static psa_status_t psa_pbkdf2_set_input_cost(
     if (pbkdf2->state != PSA_PBKDF2_STATE_INIT) {
         return PSA_ERROR_BAD_STATE;
     }
-#if UINT_MAX > 0xFFFFFFFF
+
     if (data > 0xFFFFFFFF) {
-        return PSA_ERROR_INVALID_ARGUMENT;
+        return PSA_ERROR_NOT_SUPPORTED;
     }
-#endif
+
     if (data == 0) {
         return PSA_ERROR_INVALID_ARGUMENT;
     }